| """Data pipeline for the vi/ja -> en translation task. |
| |
| The model is decoder-only, so translation is framed as plain language |
| modelling over a prompt/answer sequence: |
| |
| <bos> <vi> {vietnamese source} <en> {english target} <eos> |
| |
| Only the tokens after ``<en>`` (the English target and the closing ``<eos>``) |
| contribute to the loss. Conditioning tokens are still *read* by the model -- |
| and still get routed through the MoE, which is what makes the per-language |
| expert heatmap meaningful -- they are simply not predicted. |
| |
| Examples are packed back-to-back into fixed ``n_ctx``-token blocks rather than |
| padded individually. A typical pair is ~60 tokens, so per-example padding to |
| 256 would have thrown away three quarters of every batch, and "same number of |
| training tokens" would then have meant "same number of padding tokens". |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
|
|
| from ..utils import RUNS, load_tokenizer, train_bpe |
|
|
| DATASET = "belumind/en-vi-ja-curated-500k-triplets" |
|
|
| PAD, BOS, EOS, UNK, VI, JA, EN = 0, 1, 2, 3, 4, 5, 6 |
| LANG_VI, LANG_JA, LANG_EN, LANG_SPECIAL = 0, 1, 2, 3 |
| LANG_NAMES = ["vi (source)", "ja (source)", "en (target)", "special"] |
|
|
|
|
| def load_raw(split: str, limit: int | None = None): |
| from datasets import load_dataset |
| ds = load_dataset(DATASET, split=split) |
| if limit is not None and limit < len(ds): |
| ds = ds.select(range(limit)) |
| return ds |
|
|
|
|
| |
| |
| |
| def get_tokenizer(vocab_size: int = 16000, n_train_rows: int = 250_000, |
| force: bool = False): |
| """Train (or load) the shared byte-level BPE over all three languages.""" |
| path = RUNS / "part1" / f"tokenizer_{vocab_size}.json" |
| if path.exists() and not force: |
| return load_tokenizer(path) |
|
|
| ds = load_raw("train", limit=n_train_rows) |
|
|
| def it(): |
| |
| |
| for row in ds: |
| yield row["en"] |
| yield row["vi"] |
| yield row["ja"] |
|
|
| print(f"[part1] training BPE (vocab={vocab_size}) on {3*len(ds):,} lines") |
| return train_bpe(it(), vocab_size, path) |
|
|
|
|
| |
| |
| |
| def _encode_split(tok, split: str, limit: int | None): |
| """Return (ids, is_target, lang) int arrays for one split, packed flat.""" |
| ds = load_raw(split, limit=limit) |
| ids, tgt, lang = [], [], [] |
|
|
| en_all = tok.encode_batch([r["en"] for r in ds]) |
| vi_all = tok.encode_batch([r["vi"] for r in ds]) |
| ja_all = tok.encode_batch([r["ja"] for r in ds]) |
|
|
| for en, vi, ja in zip(en_all, vi_all, ja_all): |
| e = en.ids |
| for src_ids, marker, lang_id in ((vi.ids, VI, LANG_VI), |
| (ja.ids, JA, LANG_JA)): |
| seq = [BOS, marker] + src_ids + [EN] + e + [EOS] |
| is_t = ([0, 0] + [0] * len(src_ids) + [0] |
| + [1] * len(e) + [1]) |
| lg = ([LANG_SPECIAL, LANG_SPECIAL] + [lang_id] * len(src_ids) |
| + [LANG_SPECIAL] + [LANG_EN] * len(e) + [LANG_SPECIAL]) |
| ids.extend(seq) |
| tgt.extend(is_t) |
| lang.extend(lg) |
|
|
| return (np.asarray(ids, dtype=np.uint16), |
| np.asarray(tgt, dtype=np.uint8), |
| np.asarray(lang, dtype=np.uint8)) |
|
|
|
|
| def get_streams(tok, split: str, limit: int | None = None, force: bool = False): |
| """Tokenised streams for a split, cached on disk as .npy.""" |
| tag = f"{split}_{limit if limit is not None else 'all'}" |
| cache = RUNS / "part1" / "cache" |
| cache.mkdir(parents=True, exist_ok=True) |
| files = {k: cache / f"{tag}_{k}.npy" for k in ("ids", "tgt", "lang")} |
|
|
| if all(f.exists() for f in files.values()) and not force: |
| return tuple(np.load(f) for f in files.values()) |
|
|
| print(f"[part1] encoding split={split} limit={limit}") |
| arrs = _encode_split(tok, split, limit) |
| for f, a in zip(files.values(), arrs): |
| np.save(f, a) |
| return arrs |
|
|
|
|
| class PackedBlocks(torch.utils.data.Dataset): |
| """Chops a flat token stream into (x, y, loss_mask, lang) blocks.""" |
|
|
| def __init__(self, ids, tgt, lang, n_ctx: int): |
| n = (len(ids) - 1) // n_ctx |
| self.n_ctx, self.n = n_ctx, n |
| self.ids, self.tgt, self.lang = ids, tgt, lang |
|
|
| def __len__(self): |
| return self.n |
|
|
| def __getitem__(self, i): |
| s = i * self.n_ctx |
| e = s + self.n_ctx |
| x = torch.from_numpy(self.ids[s:e].astype(np.int64)) |
| y = torch.from_numpy(self.ids[s + 1:e + 1].astype(np.int64)) |
| m = torch.from_numpy(self.tgt[s + 1:e + 1].astype(np.float32)) |
| lg = torch.from_numpy(self.lang[s:e].astype(np.int64)) |
| return x, y, m, lg |
|
|
|
|
| def make_loader(dataset, batch_size, shuffle, seed=1337, num_workers=2): |
| g = torch.Generator().manual_seed(seed) |
| return torch.utils.data.DataLoader( |
| dataset, 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 eval_pairs(split: str = "test", n: int = 500, seed: int = 0): |
| """(lang_marker, source_text, reference_english) triples for generation.""" |
| ds = load_raw(split) |
| rng = np.random.default_rng(seed) |
| idx = rng.choice(len(ds), size=min(n, len(ds)), replace=False) |
| out = [] |
| for i in idx: |
| row = ds[int(i)] |
| out.append(("vi", row["vi"], row["en"])) |
| out.append(("ja", row["ja"], row["en"])) |
| return out |
|
|
|
|
| def build_prompt(tok, lang: str, src_text: str, max_src: int = 128): |
| marker = VI if lang == "vi" else JA |
| src = tok.encode(src_text).ids[:max_src] |
| return [BOS, marker] + src + [EN] |
|
|
