Upload src/xscript/tok/train.py with huggingface_hub

src/xscript/tok/train.py

@@ -0,0 +1,215 @@
+"""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>"]  # ids 0..3 in every flavor
+PA_REPO = "swiss-ai/parity-aware-bpe"
+
+
+# --------------------------------------------------------------------------- #
+# unigram (SentencePiece)
+# --------------------------------------------------------------------------- #
+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)   # not a TrainerSpec field in >=0.2
+    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
+
+
+# --------------------------------------------------------------------------- #
+# byte-level BPE + parity-aware BPE (swiss-ai/parity-aware-bpe)
+# --------------------------------------------------------------------------- #
+def _n_merges() -> int:
+    # vocab = 4 specials + 256 byte-level base alphabet + merges
+    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,                       # ids 0..3, in order
+        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),  # full 256 bytes
+        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")            # one file per study language
+    dev = _write_pa_dev(inputs)                    # aligned FLORES+ dev per lang
+    out = ensure(tokenizer_dir(tok_name("pa", condition)))
+    merges = out / "merges.raw.txt"
+    # parity-aware's multi-worker vocab builder is broken in the released
+    # version (pickle.load on a text-mode file), so force single-worker.
+    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)}     # 0..3
+    for ch in pre_tokenizers.ByteLevel.alphabet():                     # 256 bytes
+        vocab.setdefault(ch, len(vocab))
+    keep = max(0, VOCAB_SIZE - len(vocab))                             # merges budget
+    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:      # order guarantees this won't hit
+            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))
+    # EXACT pre-tokenizer/decoder the repo trains and loads with (byte-level)
+    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)")