Add char BPE incomplete coverage fallback

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

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

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