File size: 4,179 Bytes
5407bdc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """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
|