#!/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()