diff --git a/bindings/python/examples/train_parity_bpe.py b/bindings/python/examples/train_parity_bpe.py new file mode 100644 index 0000000000..05c902629a --- /dev/null +++ b/bindings/python/examples/train_parity_bpe.py @@ -0,0 +1,127 @@ +"""Train a parity-aware BPE tokenizer over several languages at once. + +`ParityBpeTrainer` takes one iterator per language and balances merges across them, so a +low-resource language is not left far worse compressed than a high-resource one. It needs +a per-language signal to do that, supplied either as a parallel dev set (`dev_iterators=`) +or as target compression rates (`ratio=`). Both are shown below. + +Per-language corpora are not something you are likely to have on disk, so the data comes +from the Hub: training text from `wikimedia/wikipedia`, and a parallel dev set from +`openlanguagedata/flores_plus` (997 sentences per language, aligned across languages). + +FLORES+ is gated: accept the terms at +https://huggingface.co/datasets/openlanguagedata/flores_plus and run `huggingface-cli login` +before running the dev-set demo. The ratio demo needs neither. +""" + +import datasets + +from tokenizers import Tokenizer, decoders, models, pre_tokenizers +from tokenizers.trainers import ParityBpeTrainer + + +WIKIPEDIA = "wikimedia/wikipedia" +FLORES_PLUS = "openlanguagedata/flores_plus" + +# Wikipedia uses ISO 639-1 codes, FLORES+ uses ISO 639-3 plus a script tag, so each +# language needs both: (display name, wikipedia config, flores+ config, sample sentence). +LANGUAGES = [ + ("English", "20231101.en", "eng_Latn", "The quick brown fox jumps over the lazy dog."), + ("German", "20231101.de", "deu_Latn", "Der schnelle braune Fuchs springt über den faulen Hund."), + ("French", "20231101.fr", "fra_Latn", "Le rapide renard brun saute par-dessus le chien paresseux."), + ("Hindi", "20231101.hi", "hin_Deva", "तेज़ भूरी लोमड़ी आलसी कुत्ते के ऊपर से कूदती है।"), +] + +# Target compression rates, one per language, in the order of LANGUAGES. Only the values +# relative to each other matter: the trainer picks whichever language has the lowest +# compression_rate / ratio, so a higher ratio earns more merges. +# +# ByteLevel counts compression in bytes, so equal ratios would not give equal tokenization. +# Mean bytes per FLORES+ sentence are English 125.7, German 149.4, French 155.0, Hindi +# 323.2, because Devanagari is 3 bytes per character in UTF-8 while its character count is +# close to English. These ratios are those byte counts normalized by English, which brings +# the per-language token counts closer together (in one run, Hindi went from 94.2 to 82.0 +# tokens per sentence and the spread between best and worst from 2.56x to 1.75x). +RATIOS = [1.00, 1.19, 1.23, 2.57] + +# Wikipedia articles average around 20 KB, and `train_from_iterator` materializes each +# language's corpus in memory before training, so peak memory scales with these caps. +# The defaults keep the demo near 20 MB of text. Raise both for a tokenizer you intend to use. +ARTICLES_PER_LANGUAGE = 200 +NUM_MERGES = 8000 + + +def wikipedia_iterator(config, limit=ARTICLES_PER_LANGUAGE): + """Stream up to `limit` articles for one language, so no full dump is downloaded.""" + dataset = datasets.load_dataset(WIKIPEDIA, config, split="train", streaming=True) + for count, article in enumerate(dataset): + if count >= limit: + return + yield article["text"] + + +def flores_iterator(config): + """The FLORES+ dev split: 997 sentences, the same ones in every language.""" + dataset = datasets.load_dataset(FLORES_PLUS, config, split="dev") + yield from dataset["text"] + + +def new_tokenizer(): + """Byte-level BPE, so every language is representable and compression is counted in bytes.""" + tokenizer = Tokenizer(models.BPE()) + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True) + tokenizer.decoder = decoders.ByteLevel() + return tokenizer + + +def new_trainer(): + return ParityBpeTrainer( + num_merges=NUM_MERGES, + variant="window", + window_size=100, + alpha=2.0, + initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), + ) + + +def report(tokenizer, label): + """Token count per language on one sentence, so the balance across languages is visible.""" + print(f"\n{label}: vocabulary {tokenizer.get_vocab_size()}") + for name, _, _, sentence in LANGUAGES: + print(f" {name:8s} {len(tokenizer.encode(sentence).tokens):3d} tokens") + + +# Mode 1: balance against a parallel dev set. +dev_set_tokenizer = new_tokenizer() +new_trainer().train_from_iterator( + dev_set_tokenizer, + train_iterators=[wikipedia_iterator(wiki) for _, wiki, _, _ in LANGUAGES], + dev_iterators=[flores_iterator(flores) for _, _, flores, _ in LANGUAGES], +) +dev_set_tokenizer.save("parity_bpe_devset.json") +report(dev_set_tokenizer, "dev-set balancing") + +# Mode 2: balance against target compression rates. No parallel data is involved, and the +# training iterators are rebuilt because generators are single-use. +ratio_tokenizer = new_tokenizer() +new_trainer().train_from_iterator( + ratio_tokenizer, + train_iterators=[wikipedia_iterator(wiki) for _, wiki, _, _ in LANGUAGES], + ratio=RATIOS, +) +ratio_tokenizer.save("parity_bpe_ratio.json") +report(ratio_tokenizer, "ratio balancing") + +# Neither mode reaches equal token counts, which is expected. Merges go into one shared +# vocabulary, so how well a language compresses depends on which other languages are in +# the mixture: English, German and French share the Latin script and much substring +# structure, so a merge selected for German often helps the other two, while Hindi shares +# nothing with them and has to fund its whole subword inventory from merges selected for +# Hindi. The byte budget hides a second imbalance: the four corpora are within 20% of each +# other in bytes, but Devanagari averages 2.6 bytes per character, so Hindi contributes +# about 2.2M characters of running text against English's 6.2M. `ratio` also targets a rate +# measured on the training data while parity is judged on held-out parallel text, and the +# window variant lets no language hold more than alpha / n_languages of the window (half, +# here), so a language cannot be given an unbroken run of merges to catch up. Raising +# NUM_MERGES and ARTICLES_PER_LANGUAGE closes the remaining gap more reliably than ratio +# tuning. diff --git a/bindings/python/src/trainers.rs b/bindings/python/src/trainers.rs index 0e6cd7fa05..d413a3ae21 100644 --- a/bindings/python/src/trainers.rs +++ b/bindings/python/src/trainers.rs @@ -1056,7 +1056,13 @@ where /// balances merge operations across languages using a development set or target /// compression ratios. The single training entry point is /// :meth:`train_from_iterator`, the multi-corpus analogue of -/// :meth:`tokenizers.Tokenizer.train_from_iterator`. +/// :meth:`tokenizers.Tokenizer.train_from_iterator`. Balancing requires either +/// ``dev_iterators`` (ideally parallel across languages) or ``ratio``; given +/// neither, selection follows the per-language training totals, which the +/// highest-data language dominates, so the result stays close to plain BPE. When +/// both are given the dev set takes precedence. A runnable version that obtains +/// per-language corpora and a parallel dev set is in +/// ``examples/train_parity_bpe.py``. /// /// Args: /// num_merges (:obj:`int`, `optional`): @@ -1068,6 +1074,17 @@ where /// min_frequency (:obj:`int`, `optional`): /// Minimum pair frequency to merge. Defaults to ``0``. /// +/// ratio (:obj:`List[float]`, `optional`): +/// Target compression rate per language, one entry per training iterator. The +/// trainer selects the language with the lowest ``compression_rate / ratio``, +/// so raising one language's ratio gives it more merges; only the values +/// relative to each other matter. Rates are counted in the units the +/// pre-tokenizer emits, bytes under +/// :class:`~tokenizers.pre_tokenizers.ByteLevel`, so equal ratios do not give +/// equal tokenization across scripts: Devanagari takes about 2.5x the bytes of +/// Latin script for the same content. Set each ratio proportional to the +/// language's average length on parallel text instead. Defaults to ``None``. +/// /// global_merges (:obj:`int`, `optional`): /// Number of initial standard BPE merges before switching to parity mode. Defaults to ``0``. /// @@ -1082,25 +1099,24 @@ where /// /// Example:: /// -/// from tokenizers import Tokenizer -/// from tokenizers.models import BPE -/// from tokenizers import pre_tokenizers -/// from tokenizers.trainers import ParityBpeTrainer -/// -/// tokenizer = Tokenizer(BPE()) -/// tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() -/// -/// def lines(path): -/// with open(path) as f: -/// yield from f -/// -/// trainer = ParityBpeTrainer(num_merges=32000, variant="base") -/// trainer.train_from_iterator( -/// tokenizer, -/// train_iterators=[lines("train_en.txt"), lines("train_de.txt")], -/// dev_iterators=[lines("dev_en.txt"), lines("dev_de.txt")], -/// ) -/// output = tokenizer.encode("Hello world") +/// >>> from tokenizers import Tokenizer, pre_tokenizers +/// >>> from tokenizers.models import BPE +/// >>> from tokenizers.trainers import ParityBpeTrainer +/// >>> tokenizer = Tokenizer(BPE()) +/// >>> tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel() +/// >>> trainer = ParityBpeTrainer(num_merges=32000, variant="base") +/// >>> # balance against a parallel dev set, the same sentences in every language +/// >>> trainer.train_from_iterator( +/// ... tokenizer, +/// ... train_iterators=[english_lines, hindi_lines], +/// ... dev_iterators=[english_dev, hindi_dev], +/// ... ) +/// >>> # or against target rates, which need no dev data +/// >>> trainer.train_from_iterator( +/// ... tokenizer, +/// ... train_iterators=[english_lines, hindi_lines], +/// ... ratio=[1.0, 2.57], +/// ... ) /// #[cfg(feature = "parity-aware-bpe")] #[pyclass(module = "tokenizers.trainers", name = "ParityBpeTrainer")] @@ -1312,8 +1328,9 @@ impl PyParityBpeTrainer { /// ``train_iterators``. /// /// ratio (:obj:`List[float]`, `optional`): - /// Target compression ratios per language (alternative to - /// ``dev_iterators``). + /// Target compression rates per language, an alternative to + /// ``dev_iterators`` and ignored when one is supplied. See the class + /// docstring for how to choose the values. #[pyo3(signature = (tokenizer, train_iterators, dev_iterators = None, ratio = None))] fn train_from_iterator( &self, diff --git a/docs/source-doc-builder/api/trainers.mdx b/docs/source-doc-builder/api/trainers.mdx index 0e94749442..2beade93a0 100644 --- a/docs/source-doc-builder/api/trainers.mdx +++ b/docs/source-doc-builder/api/trainers.mdx @@ -6,6 +6,10 @@ [[autodoc]] tokenizers.trainers.BpeTrainer +## ParityBpeTrainer + +[[autodoc]] tokenizers.trainers.ParityBpeTrainer + ## UnigramTrainer [[autodoc]] tokenizers.trainers.UnigramTrainer