MahmoodLab/Patho-Bench

1import os2import datasets3from datasets import Features, Value4from huggingface_hub import snapshot_download5import glob6import yaml7 8 9class PathoBenchConfig(datasets.BuilderConfig):10    def __init__(self, **kwargs):11        12        # Extract task_in_dataset and dataset_to_download from kwargs13        self.task_in_dataset = kwargs.pop("task_in_dataset", None)14        self.dataset_to_download = kwargs.pop("dataset_to_download", None)15        self.force_download = kwargs.pop("force_download", True)16        17        # Set default values for task_in_dataset and dataset_to_download18        if self.dataset_to_download is None and self.task_in_dataset is None:19            # If neither are provided, default both to '*'20            self.dataset_to_download = '*'21            self.task_in_dataset = '*'22        elif self.dataset_to_download is None and self.task_in_dataset is not None:23            # If task_in_dataset is provided but dataset_to_download is not, raise an error24            raise AssertionError("Dataset needs to be defined for the task_in_dataset provided.")25        elif self.dataset_to_download is not None and self.task_in_dataset is None:26            # If dataset_to_download is provided but task_in_dataset is not, default task_in_dataset to '*'27            self.task_in_dataset = '*'28            29        super().__init__(**kwargs)30 31 32class PathoBenchDataset(datasets.GeneratorBasedBuilder):33    """34    Downloads only the .tsv and .yaml files needed to construct the dataset.35    Excludes .png images so they don't break the builder.36    """37    BUILDER_CONFIGS = [38        PathoBenchConfig(name="custom_config", version="1.0.0", description="PathoBench config")39    ]40    BUILDER_CONFIG_CLASS = PathoBenchConfig41 42    def _info(self):43        return datasets.DatasetInfo(44            description="PathoBench: collection of canonical computational pathology tasks",45            homepage="https://github.com/mahmoodlab/patho-bench",46            license="CC BY-NC-SA 4.0 Deed",47            features=Features({48                'path': Value('string')49            })50        )51 52    def _split_generators(self, dl_manager):53        repo_id = "MahmoodLab/patho-bench"54        dataset_to_download = self.config.dataset_to_download55        local_dir = self._cache_dir_root56        force_download = self.config.force_download57        task_in_dataset = self.config.task_in_dataset58 59        # Ensure the base local directory exists60        os.makedirs(local_dir, exist_ok=True)61 62        # 1) Download the top-level available_splits.yaml63        snapshot_download(64            repo_id=repo_id,65            allow_patterns=["available_splits.yaml"],  # only this file66            repo_type="dataset",67            local_dir=local_dir,68            force_download=force_download,69        )70 71        # Read available splits72        with open(os.path.join(local_dir, "available_splits.yaml"), 'r') as file:73            available_splits = yaml.safe_load(file)74        75        # Basic validation76        if dataset_to_download != "*":77            assert dataset_to_download in available_splits, (78                f"{dataset_to_download} was not found. "79                f"Available splits: {list(available_splits.keys())}"80            )81            if task_in_dataset != "*":82                assert task_in_dataset in available_splits[dataset_to_download], (83                    f"{task_in_dataset} was not found in {dataset_to_download}. "84                    f"Available tasks: {available_splits[dataset_to_download]}"85                )86 87        # 2) Decide what to allow based on dataset/task88        #89        # We only want .tsv and the relevant .yaml files (like about.yaml, config.yaml).90        # That way, we skip .png images which can cause issues or be large in LFS.91        if dataset_to_download == "*":92            # Download every dataset subfolder's .tsv and about.yaml/config.yaml93            allow_patterns = [94                "**/*.tsv",             # All tsv splits95                "**/about.yaml",        # The about files96                "**/config.yaml",       # The config files97                "available_splits.yaml" # Already downloaded, but no harm98            ]99        else:100            if task_in_dataset == "*":101                allow_patterns = [102                    f"{dataset_to_download}/**/*.tsv",103                    f"{dataset_to_download}/**/about.yaml",104                    f"{dataset_to_download}/**/config.yaml",105                    "available_splits.yaml"106                ]107            else:108                allow_patterns = [109                    f"{dataset_to_download}/{task_in_dataset}/*.tsv",110                    f"{dataset_to_download}/{task_in_dataset}/config.yaml",111                    f"{dataset_to_download}/about.yaml",112                    "available_splits.yaml"113                ]114 115        # 3) Download the requested patterns116        snapshot_download(117            repo_id=repo_id,118            allow_patterns=allow_patterns,119            repo_type="dataset",120            local_dir=local_dir,121            force_download=force_download,122        )123        124        # 4) Locate all .tsv files to pass to _generate_examples125        search_pattern = os.path.join(local_dir, '**', '*.tsv')126        all_tsv_splits = glob.glob(search_pattern, recursive=True)127        128        return [129            datasets.SplitGenerator(130                name="full",131                gen_kwargs={"filepath": all_tsv_splits},132            )133        ]134 135    def _generate_examples(self, filepath):136        idx = 0137        for file in filepath:138            yield idx, {139                'path': file140            }141            idx += 1