From 175d46abfc2d0ae29158e33a0067413fa3cdb5b1 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Wed, 15 Jul 2026 03:42:39 -0700 Subject: [PATCH] Add char BPE incomplete coverage fallback --- data/template/nanogpt_tokenizers.py | 39 ++++++++++++++++++++++++++--- data/template/prepare.py | 2 ++ data/template/tests.py | 33 ++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 3 deletions(-) diff --git a/data/template/nanogpt_tokenizers.py b/data/template/nanogpt_tokenizers.py index 37677bd68a..17bba94d42 100644 --- a/data/template/nanogpt_tokenizers.py +++ b/data/template/nanogpt_tokenizers.py @@ -521,6 +521,7 @@ class CharBPETokenizerWithByteFallback(Tokenizer): def __init__(self, args, train_data, val_data=None): super().__init__(args) self.reuse_meta_path = getattr(args, "char_bpe_vocab_path", None) + self.incomplete_coverage_uses_bpe = getattr(args, "char_bpe_incomplete_coverage_uses_bpe", True) if self.reuse_meta_path: meta = self._load_char_bpe_meta(self.reuse_meta_path) self.desired_vocab_size = meta["vocab_size"] @@ -542,11 +543,30 @@ def __init__(self, args, train_data, val_data=None): if val_data: corpus_text += val_data - self.unique_chars = sorted(set(corpus_text)) + char_counts = Counter(corpus_text) + self.unique_chars = sorted(char_counts) if not self.unique_chars: raise ValueError("Training data must contain at least one character for char_bpe tokenization.") - self.char_tokens = list(self.unique_chars) + max_char_tokens = self.desired_vocab_size - 256 + if len(self.unique_chars) > max_char_tokens: + if not self.incomplete_coverage_uses_bpe: + raise ValueError( + "char_bpe cannot provide complete character coverage with " + f"vocab_size={self.desired_vocab_size}: found {len(self.unique_chars)} " + f"unique characters but only {max_char_tokens} non-byte token slots are available. " + "Increase vocab_size or enable char_bpe_incomplete_coverage_uses_bpe." + ) + # Keep the most frequent characters in the explicit BPE vocabulary and + # let raw byte fallback cover the rest. Ties use the character value + # for deterministic vocabulary generation. + ranked_chars = sorted(char_counts.items(), key=lambda item: (-item[1], item[0])) + self.char_tokens = [ch for ch, _count in ranked_chars[:max_char_tokens]] + self.incomplete_coverage = True + else: + self.char_tokens = list(self.unique_chars) + self.incomplete_coverage = False + self._train_merges(corpus_text) self._build_vocab() @@ -563,7 +583,7 @@ def _load_char_bpe_meta(meta_path): return meta def _train_merges(self, text): - tokens = list(text) + tokens = self._initial_bpe_tokens(text) # Nothing to merge if text empty or target vocab already satisfied if len(tokens) < 2: return @@ -602,6 +622,16 @@ def _train_merges(self, text): self.sorted_char_tokens = sorted(self.char_tokens, key=lambda t: len(t), reverse=True) + def _initial_bpe_tokens(self, text): + tokens = [] + known_chars = set(self.char_tokens) + for ch in text: + if ch in known_chars: + tokens.append(ch) + else: + tokens.extend(bytes([b]) for b in ch.encode('utf-8')) + return tokens + @staticmethod def _apply_merge(tokens, pair, new_token): merged = [] @@ -684,6 +714,9 @@ def tokenize(self, data): "char_tokens": self.char_tokens, "char_tokens_sorted": self.sorted_char_tokens, "byte_fallback": True, + "incomplete_coverage_uses_bpe": self.incomplete_coverage_uses_bpe, + "incomplete_coverage": getattr(self, "incomplete_coverage", False), + "unique_char_count": len(getattr(self, "unique_chars", self.char_tokens)), } self.finalize_meta(meta) return ids diff --git a/data/template/prepare.py b/data/template/prepare.py index d0a260f4a7..7929566c23 100644 --- a/data/template/prepare.py +++ b/data/template/prepare.py @@ -113,6 +113,8 @@ def parse_arguments(): parser.add_argument("--reuse_chars", action="store_true", help="Reuse character list from meta.pkl") parser.add_argument("--char_bpe_vocab_path", type=str, default=None, help="Path to a char_bpe meta.pkl to reuse its vocabulary/merges") + parser.add_argument("--char_bpe_incomplete_coverage_uses_bpe", action=argparse.BooleanOptionalAction, default=True, + help="When char_bpe vocab_size cannot fit every observed character plus byte fallback, keep the highest-frequency characters/BPE tokens and rely on byte fallback for the rest (default: true). Use --no-char_bpe_incomplete_coverage_uses_bpe to require complete character coverage.") # Custom tokenizer arguments parser.add_argument("--tokens_file", type=str, default=None, help="Path to the file containing newline-separated tokens for tokenization") diff --git a/data/template/tests.py b/data/template/tests.py index af3bb5baf6..72a3d3e212 100644 --- a/data/template/tests.py +++ b/data/template/tests.py @@ -310,6 +310,39 @@ def test_char_bpe_tokenizer_reuse_vocab(self): if os.path.exists(reuse_meta_path): os.remove(reuse_meta_path) + + def test_char_bpe_incomplete_coverage_uses_byte_fallback(self): + corpus = "aaaaabbbbcccdde🙂🙃" + args = Namespace(vocab_size=260, track_token_counts=True) + tokenizer = CharBPETokenizerWithByteFallback(args, corpus, None) + + self.assertEqual(tokenizer.vocab_size, 260) + self.assertEqual(len(tokenizer.char_tokens), 4) + self.assertEqual(tokenizer.char_tokens, ["a", "b", "c", "d"]) + + ids = tokenizer.tokenize(corpus) + detokenized = tokenizer.detokenize(ids) + self.assertEqual(corpus, detokenized) + self.assertTrue(any(token_id < 256 for token_id in ids)) + + with open("meta.pkl", "rb") as f: + meta = pickle.load(f) + self.assertTrue(meta["incomplete_coverage"]) + self.assertTrue(meta["incomplete_coverage_uses_bpe"]) + self.assertEqual(meta["vocab_size"], 260) + self.assertEqual(meta["unique_char_count"], len(set(corpus))) + + def test_char_bpe_incomplete_coverage_can_require_complete_coverage(self): + corpus = "abcde" + args = Namespace( + vocab_size=260, + track_token_counts=False, + char_bpe_incomplete_coverage_uses_bpe=False, + ) + + with self.assertRaisesRegex(ValueError, "complete character coverage"): + CharBPETokenizerWithByteFallback(args, corpus, None) + def test_custom_char_tokenizer_with_byte_fallback(self): args = Namespace(custom_chars_file="custom_chars.txt") # Create a custom characters file for testing