Load wikitext2 / alpaca / pg19 as HF datasets, normalised to a single 'text' column.
(name: str)
| 64 | |
| 65 | |
| 66 | def get_dataset(name: str): |
| 67 | """Load wikitext2 / alpaca / pg19 as HF datasets, normalised to a single 'text' column.""" |
| 68 | import datasets |
| 69 | if name == "wikitext2": |
| 70 | return datasets.load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1") |
| 71 | if name == "alpaca": |
| 72 | ds = datasets.load_dataset("tatsu-lab/alpaca") |
| 73 | ds = ds.remove_columns(["input", "output", "instruction"]) |
| 74 | first = ds["train"].train_test_split(test_size=0.2, seed=42) |
| 75 | second = first["test"].train_test_split(test_size=0.5, seed=42) |
| 76 | return datasets.DatasetDict({ |
| 77 | "train": first["train"], "test": second["train"], "validation": second["test"], |
| 78 | }) |
| 79 | if name == "pg19": |
| 80 | ds = datasets.load_dataset("emozilla/pg19-test", split="test") |
| 81 | ds = ds.remove_columns([c for c in ds.column_names if c != "text"]) |
| 82 | return datasets.DatasetDict({"train": ds, "test": ds, "validation": ds}) |
| 83 | raise ValueError(f"unsupported dataset: {name!r}") |
| 84 | |
| 85 | |
| 86 | def prepare_calibration_inputs( |
no outgoing calls
no test coverage detected