| """Train the study's tokenizers. |
| |
| We pretrain models with TWO SentencePiece Unigram tokenizers only. The `bpe` |
| and `pa` flavors are trained purely as tokenizer-analysis comparators for the |
| fertility/allocation gate (`xscript tok-analyze`); no model run ever uses them. |
| |
| MODEL-TRAINING tokenizers -- SentencePiece Unigram, character_coverage=0.999995, |
| byte fallback: |
| unigram_starved -- ATLAS-style replication arm: T=100 temperature mixture |
| over ~419 languages. Matches both ATLAS's ~uniform 420- |
| language mixture and the Unigram algorithm of the MADLAD- |
| 400 lineage its tokenizer descends from. |
| unigram_destarved -- the intervention arm: our 5 study languages only, byte- |
| premium content-aligned (equal *content*, not bytes, per |
| language; see data/tokcorpus.py). Same algorithm as the |
| starved arm, so the starved-vs-destarved contrast isolates |
| vocabulary allocation rather than confounding it with the |
| tokenizer algorithm. |
| |
| On the algorithm choice: Unigram is MADLAD-400's confirmed algorithm (its |
| released 256k *model* tokenizer is SentencePiece Unigram). ATLAS's 64k tokenizer |
| is a SEPARATE artifact -- trained by the MADLAD-400 authors (Kudugunta et al.) |
| on the same T=100 recipe -- whose algorithm ATLAS does not state in-text, though |
| Unigram is the natural inference from that lineage. Do not conflate ATLAS's 64k |
| with MADLAD's 256k; they are different tokenizers. |
| |
| ANALYSIS-ONLY comparators -- trained for the gate, never used to pretrain: |
| bpe -- byte-level BPE (Whitespace + ByteLevel pre-tokenization) trained with |
| HuggingFace `tokenizers`' Rust `BpeTrainer`. Quantifies how much the |
| Unigram-vs-BPE algorithm choice alone moves fertility/allocation. |
| pa -- parity-aware byte-level BPE via swiss-ai/parity-aware-bpe's |
| `parity_aware_learn_bpe.py` (window variant, for ZH), fertility- |
| equalized over the 5-way-parallel FLORES+ dev set. Same byte-level |
| alphabet as `bpe`; the merge criterion (parity-balanced vs frequency) |
| is the only difference -> a clean upper bound on fertility |
| equalization. Destarved only (it balances a fixed dev-language set). |
| Uses the slow single-threaded reference trainer -- tolerable only |
| because its corpus is 5 languages, not 419. |
| |
| Every flavor exposes exactly `VOCAB_SIZE` pieces with our four specials at ids |
| 0..3, so packed token ids stay uint16 and every downstream module stays flavor- |
| agnostic. VOCAB_SIZE is overridable via XSCRIPT_VOCAB for the CPU smoke test. |
| """ |
| import json |
| import os |
| import subprocess |
| from pathlib import Path |
|
|
| from ..langs import tok_name |
| from ..paths import TOK_CORPORA, tokenizer_dir, ensure |
| from ..data.tokcorpus import corpus_files |
|
|
| VOCAB_SIZE = int(os.environ.get("XSCRIPT_VOCAB", "65536")) |
| SPECIALS = ["<unk>", "<bos>", "<eos>", "<pad>"] |
| PA_REPO = "swiss-ai/parity-aware-bpe" |
|
|
|
|
| |
| |
| |
| def train_unigram(condition: str, seed: int = 42) -> Path: |
| import sentencepiece as spm |
| if hasattr(spm, "set_random_generator_seed"): |
| spm.set_random_generator_seed(seed) |
| files = corpus_files(condition) |
| out = ensure(tokenizer_dir(tok_name("unigram", condition))) |
| spm.SentencePieceTrainer.train( |
| input=",".join(str(f) for f in files), |
| model_prefix=str(out / "sp"), |
| model_type="unigram", |
| vocab_size=VOCAB_SIZE, |
| character_coverage=0.999995, |
| byte_fallback=True, |
| unk_id=0, bos_id=1, eos_id=2, pad_id=3, |
| unk_piece="<unk>", bos_piece="<bos>", eos_piece="<eos>", pad_piece="<pad>", |
| input_sentence_size=10_000_000, |
| shuffle_input_sentence=True, |
| train_extremely_large_corpus=True, |
| remove_extra_whitespaces=False, |
| num_threads=max(1, (os.cpu_count() or 8) - 2), |
| ) |
| _write_meta(out, "unigram", condition, files) |
| return out |
|
|
|
|
| |
| |
| |
| def _n_merges() -> int: |
| |
| return VOCAB_SIZE - len(SPECIALS) - 256 |
|
|
|
|
| def train_bpe(condition: str) -> Path: |
| from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers |
|
|
| files = corpus_files(condition) |
| out = ensure(tokenizer_dir(tok_name("bpe", condition))) |
|
|
| tok = Tokenizer(models.BPE(unk_token=None, fuse_unk=False)) |
| tok.pre_tokenizer = pre_tokenizers.Sequence( |
| [pre_tokenizers.Whitespace(), pre_tokenizers.ByteLevel(use_regex=False)]) |
| tok.decoder = decoders.ByteLevel() |
| trainer = trainers.BpeTrainer( |
| vocab_size=VOCAB_SIZE, |
| special_tokens=SPECIALS, |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), |
| show_progress=True, |
| ) |
| tok.train([str(f) for f in files], trainer) |
| tok.save(str(out / "tokenizer.json")) |
| _write_meta(out, "bpe", condition, files, |
| extra={"vocab_size_actual": tok.get_vocab_size(), |
| "source": "huggingface-tokenizers-bpe"}) |
| return out |
|
|
|
|
| def train_pa(condition: str = "destarved", variant: str = "window") -> Path: |
| if condition != "destarved": |
| raise ValueError("parity-aware BPE is destarved-only (see langs.tok_conditions)") |
| inputs = corpus_files("destarved") |
| dev = _write_pa_dev(inputs) |
| out = ensure(tokenizer_dir(tok_name("pa", condition))) |
| merges = out / "merges.raw.txt" |
| |
| |
| pa_workers = os.environ.get("XSCRIPT_PA_WORKERS", "1") |
| cmd = ["python", "-m", "parity_aware_bpe.parity_aware_learn_bpe", |
| "--variant", variant, "--symbols", str(_n_merges()), |
| "--num-workers", pa_workers, "--output", str(merges), |
| "--input", *[str(f) for f in inputs], |
| "--dev", *[str(f) for f in dev]] |
| _run(cmd) |
| _bytelevel_from_merges(merges, out, "pa", condition, inputs) |
| return out |
|
|
|
|
| def _write_pa_dev(inputs) -> list[Path]: |
| """FLORES+ dev text per language, in the SAME order as `inputs` (stem=code).""" |
| from .. import flores |
| d = ensure(TOK_CORPORA / "pa_dev") |
| dev = [] |
| for f in inputs: |
| code = f.stem |
| sents = list(flores.load(code, "dev").values()) |
| p = d / f"{code}.dev.txt" |
| p.write_text("\n".join(sents) + "\n", encoding="utf-8") |
| dev.append(p) |
| return dev |
|
|
|
|
| def _bytelevel_from_merges(merges_path: Path, out: Path, flavor: str, |
| condition: str, corpus_files_used) -> None: |
| """Merge rules -> HuggingFace byte-level BPE tokenizer, exactly VOCAB_SIZE.""" |
| from tokenizers import Tokenizer, models, pre_tokenizers, decoders |
|
|
| lines = [l.strip() for l in merges_path.read_text(encoding="utf-8").splitlines() |
| if l.strip()] |
| if lines and lines[0].startswith("#version"): |
| lines = lines[1:] |
|
|
| vocab: dict[str, int] = {s: i for i, s in enumerate(SPECIALS)} |
| for ch in pre_tokenizers.ByteLevel.alphabet(): |
| vocab.setdefault(ch, len(vocab)) |
| keep = max(0, VOCAB_SIZE - len(vocab)) |
| merges: list[tuple[str, str]] = [] |
| for line in lines: |
| if len(merges) >= keep: |
| break |
| a, b = line.split(" ") |
| if a not in vocab or b not in vocab: |
| continue |
| merges.append((a, b)) |
| vocab.setdefault(a + b, len(vocab)) |
|
|
| tok = Tokenizer(models.BPE(vocab=vocab, merges=merges, |
| unk_token=None, fuse_unk=False)) |
| |
| tok.pre_tokenizer = pre_tokenizers.Sequence( |
| [pre_tokenizers.Whitespace(), pre_tokenizers.ByteLevel(use_regex=False)]) |
| tok.decoder = decoders.ByteLevel() |
| tok.save(str(out / "tokenizer.json")) |
| _write_meta(out, flavor, condition, corpus_files_used, |
| extra={"vocab_size_actual": tok.get_vocab_size(), |
| "n_merges": len(merges), "source": PA_REPO}) |
|
|
|
|
| def _run(cmd, shell: bool = False) -> None: |
| print(f"[tok] $ {cmd if shell else ' '.join(cmd)}") |
| subprocess.run(cmd, shell=shell, check=True) |
|
|
|
|
| |
| def _write_meta(out: Path, flavor: str, condition: str, files, extra=None) -> None: |
| meta = { |
| "flavor": flavor, |
| "condition": condition, |
| "vocab_size": VOCAB_SIZE, |
| "specials": SPECIALS, |
| "corpus_files": [str(f) for f in files], |
| } |
| if extra: |
| meta.update(extra) |
| (out / "meta.json").write_text(json.dumps(meta, indent=2)) |
| print(f"[tok] trained {flavor}_{condition} -> {out}") |
|
|
|
|
| def train(flavor: str, condition: str) -> Path: |
| if flavor == "unigram": |
| return train_unigram(condition) |
| if flavor == "bpe": |
| return train_bpe(condition) |
| if flavor == "pa": |
| return train_pa(condition) |
| raise ValueError(f"unknown flavor {flavor!r} (want unigram|bpe|pa)") |
|
|
