Load and tokenise the dataset. The dataset is expected to have a "text" column. Each example is tokenised and truncated / padded to `context_length`. The labels are the same as the input ids (causal LM objective).
(tokenizer, cfg: dict)
| 112 | # ────────────────────────────────────────────────────────────────────────────── |
| 113 | |
| 114 | def build_dataloader(tokenizer, cfg: dict) -> DataLoader: |
| 115 | """ |
| 116 | Load and tokenise the dataset. |
| 117 | |
| 118 | The dataset is expected to have a "text" column. Each example is |
| 119 | tokenised and truncated / padded to `context_length`. The labels are |
| 120 | the same as the input ids (causal LM objective). |
| 121 | """ |
| 122 | ctx = cfg["context_length"] |
| 123 | |
| 124 | print("Loading dataset …") |
| 125 | ds = load_dataset(cfg["dataset_name"], split=cfg["dataset_split"], trust_remote_code=True) |
| 126 | |
| 127 | def tokenise(batch): |
| 128 | encoded = tokenizer( |
| 129 | batch["text"], |
| 130 | truncation=True, |
| 131 | max_length=ctx, |
| 132 | padding="max_length", |
| 133 | return_tensors=None, |
| 134 | ) |
| 135 | encoded["labels"] = encoded["input_ids"].copy() |
| 136 | return encoded |
| 137 | |
| 138 | ds = ds.map(tokenise, batched=True, remove_columns=ds.column_names, num_proc=4) |
| 139 | ds.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"]) |
| 140 | |
| 141 | return DataLoader( |
| 142 | ds, |
| 143 | batch_size=cfg["batch_size"], |
| 144 | shuffle=True, |
| 145 | num_workers=cfg["num_workers"], |
| 146 | pin_memory=True, |
| 147 | drop_last=True, |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | # ────────────────────────────────────────────────────────────────────────────── |