File size: 6,228 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """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
# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------
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():
# Interleaving the three languages keeps any one script from
# monopolising the merge budget.
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)
# ---------------------------------------------------------------------------
# Encoding into packed streams
# ---------------------------------------------------------------------------
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)
# ---------------------------------------------------------------------------
# Raw pairs, used for BLEU (needs real sentence boundaries, not packed blocks)
# ---------------------------------------------------------------------------
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]
|