| |
| |
| |
| |
| |
| |
| |
| |
| """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() |
|
|
