bigscience/promptsource

1# coding=utf-82import os3 4import datasets5import requests6 7from promptsource import DEFAULT_PROMPTSOURCE_CACHE_HOME8from promptsource.templates import INCLUDED_USERS9 10 11def removeHyphen(example):12    example_clean = {}13    for key in example.keys():14        if "-" in key:15            new_key = key.replace("-", "_")16            example_clean[new_key] = example[key]17        else:18            example_clean[key] = example[key]19    example = example_clean20    return example21 22 23def renameDatasetColumn(dataset):24    col_names = dataset.column_names25    for cols in col_names:26        if "-" in cols:27            dataset = dataset.rename_column(cols, cols.replace("-", "_"))28    return dataset29 30 31#32# Helper functions for datasets library33#34 35 36def get_dataset_builder(path, conf=None):37    "Get a dataset builder from name and conf."38    module_path = datasets.load.dataset_module_factory(path)39    builder_cls = datasets.load.import_main_class(module_path.module_path, dataset=True)40    if conf:41        builder_instance = builder_cls(name=conf, cache_dir=None, hash=module_path.hash)42    else:43        builder_instance = builder_cls(cache_dir=None, hash=module_path.hash)44    return builder_instance45 46 47def get_dataset(path, conf=None):48    "Get a dataset from name and conf."49    builder_instance = get_dataset_builder(path, conf)50    if builder_instance.manual_download_instructions is None and builder_instance.info.size_in_bytes is not None:51        builder_instance.download_and_prepare()52        return builder_instance.as_dataset()53    else:54        return load_dataset(path, conf)55 56 57def load_dataset(dataset_name, subset_name):58    try:59        return datasets.load_dataset(dataset_name, subset_name)60    except datasets.builder.ManualDownloadError:61        cache_root_dir = (62            os.environ["PROMPTSOURCE_MANUAL_DATASET_DIR"]63            if "PROMPTSOURCE_MANUAL_DATASET_DIR" in os.environ64            else DEFAULT_PROMPTSOURCE_CACHE_HOME65        )66        data_dir = (67            f"{cache_root_dir}/{dataset_name}"68            if subset_name is None69            else f"{cache_root_dir}/{dataset_name}/{subset_name}"70        )71        return datasets.load_dataset(72            dataset_name,73            subset_name,74            data_dir=data_dir,75        )76 77 78def get_dataset_confs(path):79    "Get the list of confs for a dataset."80    module_path = datasets.load.dataset_module_factory(path).module_path81    # Get dataset builder class from the processing script82    builder_cls = datasets.load.import_main_class(module_path, dataset=True)83    # Instantiate the dataset builder84    confs = builder_cls.BUILDER_CONFIGS85    if confs and len(confs) > 1:86        return confs87    return []88 89 90def render_features(features):91    """Recursively render the dataset schema (i.e. the fields)."""92    if isinstance(features, dict):93        return {k: render_features(v) for k, v in features.items()}94    if isinstance(features, datasets.features.ClassLabel):95        return features.names96 97    if isinstance(features, datasets.features.Value):98        return features.dtype99 100    if isinstance(features, datasets.features.Sequence):101        return {"[]": render_features(features.feature)}102    return features103 104 105#106# Loads dataset information107#108 109 110def filter_english_datasets():111    """112    Filter English datasets based on language tags in metadata.113 114    Also includes the datasets of any users listed in INCLUDED_USERS115    """116    english_datasets = []117 118    response = requests.get("https://huggingface.co/api/datasets?full=true")119    tags = response.json()120 121    for dataset in tags:122        dataset_name = dataset["id"]123 124        is_community_dataset = "/" in dataset_name125        if is_community_dataset:126            user = dataset_name.split("/")[0]127            if user in INCLUDED_USERS:128                english_datasets.append(dataset_name)129            continue130 131        if "cardData" not in dataset:132            continue133        metadata = dataset["cardData"]134 135        if "language" not in metadata:136            continue137        languages = metadata["language"]138 139        if "en" in languages or "en-US" in languages:140            english_datasets.append(dataset_name)141 142    return sorted(english_datasets)143 144 145def list_datasets():146    """Get all the datasets to work with."""147    dataset_list = filter_english_datasets()148    dataset_list.sort(key=lambda x: x.lower())149    return dataset_list150