| """Data pipeline for Part 2: next-token pretraining on the human-AI corpus. |
| |
| Every document in ``browndw/human-ai-parallel-corpus`` exists in several |
| parallel versions -- the human original (``@chunk_N``) and one continuation |
| per generator model (``@gpt-4o-*``). Splitting rows at random would drop |
| near-paraphrases of the same underlying document into both train and test and |
| quietly inflate every validation number, so the split is made on the *base* |
| document id and all of a document's versions move together. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
|
|
| from ..utils import RUNS, load_tokenizer, train_bpe |
|
|
| DATASET = "browndw/human-ai-parallel-corpus" |
| PAD, BOS, EOS, UNK = 0, 1, 2, 3 |
| SPECIALS = ["<pad>", "<bos>", "<eos>", "<unk>"] |
|
|
|
|
| def _base_id(doc_id: str) -> str: |
| return doc_id.split("@")[0] |
|
|
|
|
| def load_splits(val_frac=0.02, test_frac=0.02, seed=1337): |
| """Document-level train/val/test split of the corpus.""" |
| from datasets import load_dataset |
| ds = load_dataset(DATASET, split="train") |
|
|
| bases = sorted({_base_id(d) for d in ds["doc_id"]}) |
| rng = np.random.default_rng(seed) |
| rng.shuffle(bases) |
| n_val = int(len(bases) * val_frac) |
| n_test = int(len(bases) * test_frac) |
| val_b = set(bases[:n_val]) |
| test_b = set(bases[n_val:n_val + n_test]) |
|
|
| which = [] |
| for d in ds["doc_id"]: |
| b = _base_id(d) |
| which.append("val" if b in val_b else "test" if b in test_b else "train") |
| which = np.array(which) |
| idx = {s: np.flatnonzero(which == s) for s in ("train", "val", "test")} |
| return ds, idx |
|
|
|
|
| def get_tokenizer(vocab_size: int = 16000, force: bool = False): |
| path = RUNS / "part2" / f"tokenizer_{vocab_size}.json" |
| if path.exists() and not force: |
| return load_tokenizer(path) |
| ds, idx = load_splits() |
| train_texts = ds.select(idx["train"])["text"] |
| print(f"[part2] training BPE (vocab={vocab_size}) on {len(train_texts):,} docs") |
| return train_bpe(iter(train_texts), vocab_size, path, specials=SPECIALS) |
|
|
|
|
| def get_stream(tok, split: str, force: bool = False) -> np.ndarray: |
| """Flat token stream for one split, documents joined by <eos>.""" |
| cache = RUNS / "part2" / "cache" |
| cache.mkdir(parents=True, exist_ok=True) |
| f = cache / f"{split}.npy" |
| if f.exists() and not force: |
| return np.load(f) |
|
|
| ds, idx = load_splits() |
| texts = ds.select(idx[split])["text"] |
| print(f"[part2] encoding {split}: {len(texts):,} docs") |
| ids = [] |
| B = 2000 |
| for s in range(0, len(texts), B): |
| for e in tok.encode_batch(texts[s:s + B]): |
| ids.append(BOS) |
| ids.extend(e.ids) |
| ids.append(EOS) |
| arr = np.asarray(ids, dtype=np.uint16) |
| np.save(f, arr) |
| return arr |
|
|
|
|
| class Blocks(torch.utils.data.Dataset): |
| """Contiguous (x, y) blocks over a flat token stream.""" |
|
|
| def __init__(self, ids: np.ndarray, n_ctx: int): |
| self.ids, self.n_ctx = ids, n_ctx |
| self.n = (len(ids) - 1) // n_ctx |
|
|
| def __len__(self): |
| return self.n |
|
|
| def __getitem__(self, i): |
| s = i * self.n_ctx |
| x = torch.from_numpy(self.ids[s:s + self.n_ctx].astype(np.int64)) |
| y = torch.from_numpy(self.ids[s + 1:s + 1 + self.n_ctx].astype(np.int64)) |
| return x, y |
|
|
|
|
| def make_loader(ds, batch_size, shuffle, seed=1337, num_workers=2): |
| g = torch.Generator().manual_seed(seed) |
| return torch.utils.data.DataLoader( |
| ds, batch_size=batch_size, shuffle=shuffle, |
| generator=g if shuffle else None, num_workers=num_workers, |
| pin_memory=True, drop_last=True, persistent_workers=num_workers > 0) |
|
|
|
|
| def bleu_prompts(tok, n=300, prompt_len=64, cont_len=48, seed=0): |
| """(prompt_ids, reference_text) pairs from held-out test documents.""" |
| ds, idx = load_splits() |
| rng = np.random.default_rng(seed) |
| pick = rng.choice(idx["test"], size=min(n, len(idx["test"])), replace=False) |
| out = [] |
| for i in pick: |
| e = tok.encode(ds[int(i)]["text"]).ids |
| if len(e) < prompt_len + cont_len: |
| continue |
| out.append(([BOS] + e[:prompt_len], |
| tok.decode(e[prompt_len:prompt_len + cont_len]))) |
| return out |
|
|
