training_scripts/train_tokenizer.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "datasets>=4.4.2",
#   "tokenizers>=0.22.2",
# ]
# ///
"""Rebuild the tokenizer released with Bekko Static Embedding Multilingual v1.

What this script does
---------------------
This standalone entry point downloads deterministic text prefixes from the
pinned public ``hotchpotch/bekko-embedding-v1-unsupervised`` revision. It first
trains a 30,000-token English WordPiece core, preserves those tokens and IDs,
then extends the vocabulary with multilingual candidates to exactly 200,000
tokens. It validates the result and writes ``tokenizer.json`` plus a provenance
manifest. Tokenizer training uses non-bitext data only.

How to run it
-------------
The script contains PEP 723 dependency metadata, so the shortest invocation is::

    uv run train_tokenizer.py \
        --corpus-dir output/tokenizer-corpus \
        --output-dir output/tokenizer

Use persistent paths because downloading and preparing the corpus takes time.
Run ``uv run train_tokenizer.py --help`` for all options. The generated
tokenizer is standard uncased WordPiece and needs no custom inference code.

Reproducing v1
--------------
Run this script before ``train_model.py``. Then set
``model.tokenizer_model`` in the v1 YAML recipe to the tokenizer output
directory. Do not change the pinned dataset revision, corpus sizes, vocabulary
sizes, ordering, or seed when an exact v1 reproduction is required. The model
trainer keeps sequences at 512 tokens during training and changes only the
final exported tokenizer ceiling to 32,768, as specified by the YAML recipe.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import time
from pathlib import Path
from typing import Any

from datasets import load_dataset
from tokenizers import (
    Tokenizer,
    decoders,
    normalizers,
    pre_tokenizers,
    processors,
    trainers,
)
from tokenizers.models import WordPiece

DATASET_ID = "hotchpotch/bekko-embedding-v1-unsupervised"
DATASET_REVISION = "4db86946dc042f8f97b675bbdbe281287c254138"
RECIPE_ID = "bekko-multilingual-wp200k-en30k-uncased"
ALGORITHM = "wordpiece"
VOCAB_SIZE = 200000
ENGLISH_CORE_SIZE = 30000
LOWERCASE = True
SCRIPT_BALANCED = False
NORMALIZATION = "bert_uncased"
PRE_TOKENIZER = "bert"
ROWS_PER_LANGUAGE = 32_768
ENGLISH_ROWS_PER_SOURCE = 32_768
TEXT_COLUMN = "pos"
SPECIAL_TOKENS = ["[UNK]", "[PAD]", "[CLS]", "[SEP]", "[MASK]"]
ENGLISH_FINEWEB_CONFIG = "hotchpotch-fineweb-ir-simulated-search-queries__english__qd"
LANGUAGE_CONFIGS = {
    "ar": "hotchpotch-wikipedia-multilingual-ir-pairs__arwiki-20251202-v1.1.0__qd",
    "bg": "huggingfacefw-finewiki__bg__qd",
    "bn": "huggingfacefw-finewiki__bn__qd",
    "ca": "huggingfacefw-finewiki__ca__qd",
    "cs": "huggingfacefw-finewiki__cs__qd",
    "da": "huggingfacefw-finewiki__da__qd",
    "de": "hotchpotch-wikipedia-multilingual-ir-pairs__dewiki-20251202-v1.1.0__qd",
    "el": "huggingfacefw-finewiki__el__qd",
    "en": "hotchpotch-wikipedia-multilingual-ir-pairs__enwiki-20251202-v1.1.0__qd",
    "es": "hotchpotch-wikipedia-multilingual-ir-pairs__eswiki-20251202-v1.1.0__qd",
    "et": "huggingfacefw-finewiki__et__qd",
    "fa": "huggingfacefw-finewiki__fa__qd",
    "fi": "huggingfacefw-finewiki__fi__qd",
    "fr": "hotchpotch-wikipedia-multilingual-ir-pairs__frwiki-20251202-v1.1.0__qd",
    "he": "huggingfacefw-finewiki__he__qd",
    "hi": "huggingfacefw-finewiki__hi__qd",
    "hr": "huggingfacefw-finewiki__hr__qd",
    "hu": "huggingfacefw-finewiki__hu__qd",
    "id": "huggingfacefw-finewiki__id__qd",
    "it": "hotchpotch-wikipedia-multilingual-ir-pairs__itwiki-20251202-v1.1.0__qd",
    "ja": "hotchpotch-wikipedia-multilingual-ir-pairs__jawiki-20251202-v1.1.0__qd",
    "ko": "hotchpotch-wikipedia-multilingual-ir-pairs__kowiki-20251202-v1.1.0__qd",
    "nl": "huggingfacefw-finewiki__nl__qd",
    "no": "huggingfacefw-finewiki__no__qd",
    "pl": "huggingfacefw-finewiki__pl__qd",
    "pt": "hotchpotch-wikipedia-multilingual-ir-pairs__ptwiki-20251202-v1.1.0__qd",
    "ro": "huggingfacefw-finewiki__ro__qd",
    "ru": "hotchpotch-wikipedia-multilingual-ir-pairs__ruwiki-20251202-v1.1.0__qd",
    "sk": "huggingfacefw-finewiki__sk__qd",
    "sl": "huggingfacefw-finewiki__sl__qd",
    "sr": "huggingfacefw-finewiki__sr__qd",
    "sv": "huggingfacefw-finewiki__sv__qd",
    "sw": "huggingfacefw-finewiki__sw__qd",
    "te": "huggingfacefw-finewiki__te__qd",
    "th": "huggingfacefw-finewiki__th__qd",
    "tr": "huggingfacefw-finewiki__tr__qd",
    "uk": "huggingfacefw-finewiki__uk__qd",
    "vi": "huggingfacefw-finewiki__vi__qd",
    "yo": "huggingfacefw-finewiki__yo__qd",
    "zh": "hotchpotch-wikipedia-multilingual-ir-pairs__zhwiki-20251202-v1.1.0__qd",
}
SCRIPT_GROUPS = {
    "arabic": {"ar", "fa"},
    "cjk": {"ja", "ko", "zh"},
    "cyrillic": {"bg", "ru", "sr", "uk"},
    "greek": {"el"},
    "indic": {"bn", "hi", "te"},
    "hebrew": {"he"},
    "thai": {"th"},
    "latin": {"ca", "cs", "da", "de", "es", "et", "fi", "fr", "hr", "hu",
              "id", "it", "nl", "no", "pl", "pt", "ro", "sk", "sl", "sv",
              "sw", "tr", "vi", "yo"},
}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def normalize_document(value: Any) -> str:
    text = " ".join(str(value).split())
    if not text:
        raise ValueError("Bekko document is empty")
    return text


def materialize_source(path: Path, config: str, rows: int) -> dict[str, Any]:
    """Download exactly one deterministic streaming prefix, with safe resume."""
    metadata_path = path.with_suffix(".json")
    if path.is_file() and metadata_path.is_file():
        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
        if (metadata.get("rows") == rows and metadata.get("config") == config
                and metadata.get("revision") == DATASET_REVISION
                and metadata.get("sha256") == sha256_file(path)):
            return metadata
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(".txt.incomplete")
    stream = load_dataset(
        DATASET_ID, config, revision=DATASET_REVISION, split="train", streaming=True
    )
    written = 0
    with temporary.open("w", encoding="utf-8") as handle:
        for row in stream.take(rows):
            handle.write(normalize_document(row[TEXT_COLUMN]) + "\n")
            written += 1
    if written != rows:
        temporary.unlink(missing_ok=True)
        raise RuntimeError(f"{config} supplied {written:,}/{rows:,} rows")
    os.replace(temporary, path)
    metadata = {
        "config": config, "dataset": DATASET_ID, "revision": DATASET_REVISION,
        "rows": written, "sha256": sha256_file(path), "text_column": TEXT_COLUMN,
    }
    metadata_path.write_text(
        json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(f"corpus {path.stem}: {written:,} documents", flush=True)
    return metadata


def prepare_corpus(
    corpus_dir: Path, rows_per_language: int, english_rows_per_source: int
) -> tuple[list[Path], dict[str, Path]]:
    """Build 39-language extension data and a 1:1 English Wiki/FineWeb core."""
    root = corpus_dir / (
        f"multi-{rows_per_language}-english-{english_rows_per_source}-v2"
    )
    multilingual: list[Path] = []
    sources: dict[str, Any] = {}
    for language, config in sorted(LANGUAGE_CONFIGS.items()):
        if language == "en":
            continue
        path = root / "multilingual" / f"{language}.txt"
        sources[language] = materialize_source(path, config, rows_per_language)
        multilingual.append(path)
    english_paths = {
        "wikipedia": root / "english" / "wikipedia.txt",
        "fineweb": root / "english" / "fineweb.txt",
    }
    sources["en_wikipedia"] = materialize_source(
        english_paths["wikipedia"], LANGUAGE_CONFIGS["en"], english_rows_per_source
    )
    sources["en_fineweb"] = materialize_source(
        english_paths["fineweb"], ENGLISH_FINEWEB_CONFIG, english_rows_per_source
    )
    manifest = {
        "bitext_included": False,
        "english_sampling": "equal document count: Bekko Wikipedia and FineWeb",
        "english_rows_per_source": english_rows_per_source,
        "extension_language_count": 39,
        "rows_per_extension_language": rows_per_language,
        "sources": sources,
    }
    (root / "manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    return multilingual, english_paths


def build_wordpiece(vocab_size: int, lowercase: bool = LOWERCASE) -> tuple[Tokenizer, trainers.WordPieceTrainer]:
    if vocab_size <= len(SPECIAL_TOKENS):
        raise ValueError("vocab_size must exceed special-token count")
    tokenizer = Tokenizer(WordPiece(unk_token="[UNK]"))
    tokenizer.normalizer = normalizers.BertNormalizer(
        clean_text=True, handle_chinese_chars=True,
        strip_accents=lowercase, lowercase=lowercase,
    )
    tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
    tokenizer.decoder = decoders.WordPiece(prefix="##", cleanup=True)
    trainer = trainers.WordPieceTrainer(
        vocab_size=vocab_size, min_frequency=2, special_tokens=SPECIAL_TOKENS,
        continuing_subword_prefix="##", limit_alphabet=16_000,
    )
    return tokenizer, trainer


def extension_training_paths(paths: list[Path]) -> list[Path]:
    if not SCRIPT_BALANCED:
        return paths
    by_language = {path.stem: path for path in paths}
    largest = max(len(group) for group in SCRIPT_GROUPS.values())
    weighted: list[Path] = []
    for languages in SCRIPT_GROUPS.values():
        repeats = math.ceil(largest / len(languages))
        for language in sorted(languages):
            weighted.extend([by_language[language]] * repeats)
    if set(by_language) != set().union(*SCRIPT_GROUPS.values()):
        raise RuntimeError("SCRIPT_GROUPS must cover every extension language once")
    return weighted


def merge_vocabularies(
    english: Tokenizer, multilingual: Tokenizer, total_size: int
) -> tuple[Tokenizer, list[str]]:
    """Preserve every English token ID, then append multilingual candidates."""
    english_order = [
        token for token, _ in sorted(english.get_vocab().items(), key=lambda item: item[1])
    ]
    candidate_order = [
        token for token, _ in sorted(multilingual.get_vocab().items(), key=lambda item: item[1])
    ]
    merged = list(english_order)
    seen = set(merged)
    for token in candidate_order:
        if token not in seen:
            merged.append(token)
            seen.add(token)
            if len(merged) == total_size:
                break
    if len(merged) != total_size:
        raise RuntimeError(f"Only {len(merged):,}/{total_size:,} unique tokens available")
    vocab = {token: index for index, token in enumerate(merged)}
    tokenizer = Tokenizer(
        WordPiece(vocab=vocab, unk_token="[UNK]", continuing_subword_prefix="##")
    )
    tokenizer.normalizer = normalizers.BertNormalizer(
        clean_text=True, handle_chinese_chars=True,
        strip_accents=LOWERCASE, lowercase=LOWERCASE,
    )
    tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
    tokenizer.decoder = decoders.WordPiece(prefix="##", cleanup=True)
    tokenizer.post_processor = processors.TemplateProcessing(
        single="[CLS] $A [SEP]", pair="[CLS] $A [SEP] $B [SEP]",
        special_tokens=[
            ("[CLS]", vocab["[CLS]"]), ("[SEP]", vocab["[SEP]"])
        ],
    )
    return tokenizer, english_order


def train_tokenizer(
    *, corpus_dir: Path, output_dir: Path,
    rows_per_language: int = ROWS_PER_LANGUAGE,
    english_rows_per_source: int = ENGLISH_ROWS_PER_SOURCE,
    vocab_size: int = VOCAB_SIZE,
    english_core_size: int = ENGLISH_CORE_SIZE,
) -> Path:
    if english_core_size >= vocab_size:
        raise ValueError("english_core_size must be smaller than vocab_size")
    multilingual_paths, english_paths = prepare_corpus(
        corpus_dir, rows_per_language, english_rows_per_source
    )
    output_dir.mkdir(parents=True, exist_ok=True)
    started = time.monotonic()

    english, english_trainer = build_wordpiece(english_core_size)
    english.train(
        [str(english_paths["wikipedia"]), str(english_paths["fineweb"])],
        english_trainer,
    )
    if english.get_vocab_size() != english_core_size:
        raise RuntimeError("English core did not reach its exact requested size")

    candidate_size = vocab_size + english_core_size
    multilingual, multilingual_trainer = build_wordpiece(candidate_size)
    multilingual.train(
        [str(path) for path in extension_training_paths(multilingual_paths)],
        multilingual_trainer,
    )
    tokenizer, english_order = merge_vocabularies(english, multilingual, vocab_size)
    final_vocab = tokenizer.get_vocab()
    if any(final_vocab[token] != index for index, token in enumerate(english_order)):
        raise RuntimeError("English core token IDs changed during extension")
    tokenizer.enable_truncation(max_length=512)

    tokenizer_path = output_dir / "tokenizer.json"
    tokenizer.save(str(tokenizer_path))
    (output_dir / "tokenizer_config.json").write_text(json.dumps({
        "clean_up_tokenization_spaces": True,
        "do_lower_case": LOWERCASE,
        "model_max_length": 512,
        "tokenizer_class": "BertTokenizerFast",
        "unk_token": "[UNK]", "pad_token": "[PAD]", "cls_token": "[CLS]",
        "sep_token": "[SEP]", "mask_token": "[MASK]",
    }, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    (output_dir / "special_tokens_map.json").write_text(json.dumps({
        "unk_token": "[UNK]", "pad_token": "[PAD]", "cls_token": "[CLS]",
        "sep_token": "[SEP]", "mask_token": "[MASK]",
    }, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    (output_dir / "english-core-vocab.json").write_text(
        json.dumps(english_order, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    manifest = {
        "algorithm": ALGORITHM, "bitext_included": False,
        "dataset": DATASET_ID, "dataset_revision": DATASET_REVISION,
        "english_core_ids_preserved": True,
        "english_core_size": english_core_size,
        "english_sampling": "Bekko Wikipedia:FineWeb = 1:1 documents",
        "extension_language_count": 39, "lowercase": LOWERCASE,
        "recipe_id": RECIPE_ID, "rows_per_language": rows_per_language,
        "english_rows_per_source": english_rows_per_source,
        "script_balanced": SCRIPT_BALANCED,
        "runtime_seconds": time.monotonic() - started,
        "tokenizer_sha256": sha256_file(tokenizer_path),
        "vocab_size_actual": tokenizer.get_vocab_size(),
        "vocab_size_requested": vocab_size,
    }
    (output_dir / "manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    return output_dir


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--corpus-dir", type=Path, default=Path("output/multilingual-tokenizers/corpus"))
    parser.add_argument("--output-dir", type=Path, default=Path("output/multilingual-tokenizers") / RECIPE_ID)
    parser.add_argument("--rows-per-language", type=int, default=ROWS_PER_LANGUAGE)
    parser.add_argument("--english-rows-per-source", type=int, default=ENGLISH_ROWS_PER_SOURCE)
    parser.add_argument("--vocab-size", type=int, default=VOCAB_SIZE)
    parser.add_argument("--english-core-size", type=int, default=ENGLISH_CORE_SIZE)
    args = parser.parse_args()
    print(train_tokenizer(
        corpus_dir=args.corpus_dir, output_dir=args.output_dir,
        rows_per_language=args.rows_per_language,
        english_rows_per_source=args.english_rows_per_source,
        vocab_size=args.vocab_size, english_core_size=args.english_core_size,
    ))


if __name__ == "__main__":
    main()