Fixed-length PPL loader (lm-eval-harness wikitext protocol: concat-and-chunk).
(
tokenizer, dataset: str, seqlen: int, batch_size: int,
split: str = "test", max_chunks: int | None = None,
)
| 611 | |
| 612 | |
| 613 | def prepare_ppl_dataloader( |
| 614 | tokenizer, dataset: str, seqlen: int, batch_size: int, |
| 615 | split: str = "test", max_chunks: int | None = None, |
| 616 | ): |
| 617 | """Fixed-length PPL loader (lm-eval-harness wikitext protocol: concat-and-chunk).""" |
| 618 | ds = get_dataset(dataset)[split] |
| 619 | col = ds.column_names[0] |
| 620 | full = "\n\n".join(r[col] for r in ds if r[col]) |
| 621 | ids = tokenizer(full, return_tensors="pt", truncation=False).input_ids[0] |
| 622 | n_chunks = ids.numel() // seqlen |
| 623 | if n_chunks == 0: |
| 624 | raise ValueError(f"{dataset}/{split} too short for seqlen={seqlen}") |
| 625 | if max_chunks is not None: |
| 626 | n_chunks = min(n_chunks, max_chunks) |
| 627 | chunks = ids[: n_chunks * seqlen].view(n_chunks, seqlen).contiguous() |
| 628 | |
| 629 | class _DS(torch.utils.data.Dataset): |
| 630 | def __len__(self): return chunks.shape[0] |
| 631 | def __getitem__(self, i): |
| 632 | return {"input_ids": chunks[i], |
| 633 | "attention_mask": torch.ones(seqlen, dtype=torch.long)} |
| 634 | return torch.utils.data.DataLoader(_DS(), batch_size=batch_size, shuffle=False) |
| 635 | |
| 636 | |
| 637 | @torch.no_grad() |
no test coverage detected