Pack ``nsamples`` chunks of exactly ``seqlen`` tokens from random rows.
(
tokenizer, dataset: str, nsamples: int, seqlen: int, seed: int = 42, split: str = "train",
)
| 84 | |
| 85 | |
| 86 | def prepare_calibration_inputs( |
| 87 | tokenizer, dataset: str, nsamples: int, seqlen: int, seed: int = 42, split: str = "train", |
| 88 | ) -> list[torch.Tensor]: |
| 89 | """Pack ``nsamples`` chunks of exactly ``seqlen`` tokens from random rows.""" |
| 90 | ds = get_dataset(dataset)[split] |
| 91 | col = ds.column_names[0] |
| 92 | ds = ds.filter(lambda r: r[col] is not None and len(r[col]) > 0) |
| 93 | gen = torch.Generator().manual_seed(seed) |
| 94 | out: list[torch.Tensor] = [] |
| 95 | pbar = tqdm(total=nsamples, desc=f"Packing {dataset}", unit="sample") |
| 96 | while len(out) < nsamples: |
| 97 | text = "" |
| 98 | for _ in range(10000): |
| 99 | idx = int(torch.randint(0, len(ds), (1,), generator=gen).item()) |
| 100 | text = ds[idx][col] if not text else f"{text}\n\n{ds[idx][col]}" |
| 101 | ids = tokenizer(text, return_tensors="pt", add_special_tokens=False).input_ids[0] |
| 102 | if ids.numel() >= seqlen: |
| 103 | out.append(ids[:seqlen].contiguous()) |
| 104 | pbar.update(1) |
| 105 | break |
| 106 | else: |
| 107 | raise RuntimeError(f"failed to pack {seqlen} tokens from {dataset}") |
| 108 | pbar.close() |
| 109 | return out |
| 110 | |
| 111 | |
| 112 | @torch.no_grad() |