Get the dataset from the HuggingFace datasets library. Args: name: The name of the HuggingFace dataset to load. Must be one of "wikitext2", "ptb", "c4" or "alpaca". Returns: The dataset.
(name: str)
| 10 | |
| 11 | |
| 12 | def get_dataset(name: str) -> datasets.DatasetDict: |
| 13 | """ |
| 14 | Get the dataset from the HuggingFace datasets library. |
| 15 | |
| 16 | Args: |
| 17 | name: The name of the HuggingFace dataset to load. Must be one of "wikitext2", "ptb", "c4" or "alpaca". |
| 18 | |
| 19 | Returns: |
| 20 | The dataset. |
| 21 | """ |
| 22 | logging.info(f"Loading dataset: {name}") |
| 23 | |
| 24 | ds_properties = { |
| 25 | "wikitext2": {"path": "wikitext", "config_name": "wikitext-2-raw-v1"}, |
| 26 | "ptb": {"path": "ptb_text_only", "config_name": "penn_treebank"}, |
| 27 | "c4": { |
| 28 | "path": "allenai/c4", |
| 29 | "config_name": "allenai--c4", |
| 30 | "data_files": { |
| 31 | "train": "en/c4-train.00000-of-01024.json.gz", |
| 32 | "validation": "en/c4-validation.00000-of-00008.json.gz", |
| 33 | }, |
| 34 | "cols_to_remove": ['url', 'timestamp'], |
| 35 | }, |
| 36 | "alpaca": {"path": "tatsu-lab/alpaca", "cols_to_remove": ['input', 'output', 'instruction']}, |
| 37 | } |
| 38 | |
| 39 | if name not in ds_properties: |
| 40 | raise NotImplementedError("The provided dataset is not supported") |
| 41 | |
| 42 | properties = ds_properties[name] |
| 43 | ds = datasets.load_dataset( |
| 44 | properties["path"], name=properties.get("config_name"), data_files=properties.get("data_files") |
| 45 | ) |
| 46 | |
| 47 | if "cols_to_remove" in properties: |
| 48 | ds = ds.remove_columns(properties["cols_to_remove"]) |
| 49 | |
| 50 | # if alpaca, create a test and validation set from the training set |
| 51 | if name == "alpaca": |
| 52 | ds = ds["train"].train_test_split(test_size=0.2, seed=42) |
| 53 | temp_ds = ds.pop("test") |
| 54 | temp_ds = temp_ds.train_test_split(test_size=0.5, seed=42) |
| 55 | ds["test"] = temp_ds["train"] |
| 56 | ds["validation"] = temp_ds["test"] |
| 57 | |
| 58 | logging.info("Loading dataset done") |
| 59 | return ds |
| 60 | |
| 61 | |
| 62 | def prepare_test_dataloader( |