def get_dataset_fn(dataset_type):
    """
    Helper function to get the dataset function based on the dataset type
    """
    if dataset_type == "image_text":
        return get_laion_dataset
    elif dataset_type == "mmc4":
        return get_mmc4_dataset
    else:
        raise ValueError(f"Unsupported dataset type: {dataset_type}")
Helper function to get the dataset function based on the dataset type
open_flamingo/train/data.py