feat(tokenizer): exclude HuggingFace input token IDs

data/template/nanogpt_tokenizers.py

@@ -305,6 +305,11 @@ def __init__(self, args):
             self._explain_load_error(exc)
             raise
 
+        self.excluded_token_ids = sorted(set(
+            getattr(args, "hf_exclude_token_ids", None) or []
+        ))
+        self.input_tokenizer = self._build_input_tokenizer()
+
         # Try to resolve the actual commit SHA the tokenizer was loaded from,
         # for reproducibility. This is best-effort and never fatal.
         self.hf_resolved_commit = self._resolve_commit_hash()
@@ -314,8 +319,65 @@ def __init__(self, args):
         meta_output_path = getattr(args, "meta_output_path", "meta.pkl")
         output_dir = os.path.dirname(meta_output_path) or "."
         self.hf_local_dir = os.path.join(output_dir, "hf_tokenizer")
+        self.hf_input_local_dir = os.path.join(output_dir, "hf_input_tokenizer")
         self.last_token_count = 0
 
+    def _build_input_tokenizer(self):
+        """Return an encoder whose BPE graph cannot emit excluded IDs.
+
+        Merely filtering IDs after encoding corrupts the text stream. Instead we
+        remove the corresponding vocabulary nodes and every merge touching (or
+        creating) them, allowing BPE to fall back to smaller pieces. The decoder
+        intentionally remains the unmodified tokenizer, so model output retains
+        the complete vocabulary.
+        """
+        if not self.excluded_token_ids:
+            return self.tokenizer
+        if not getattr(self.tokenizer, "is_fast", False):
+            raise ValueError("--hf_exclude_token_ids requires a fast HuggingFace tokenizer")
+
+        from tokenizers import Tokenizer as BackendTokenizer
+
+        backend = getattr(self.tokenizer, "backend_tokenizer", None)
+        config = json.loads(backend.to_str())
+        model = config.get("model", {})
+        if model.get("type") != "BPE":
+            raise ValueError(
+                "--hf_exclude_token_ids currently supports BPE tokenizers only; "
+                f"loaded backend model is {model.get('type')!r}"
+            )
+
+        vocab = model.get("vocab", {})
+        id_to_piece = {int(token_id): piece for piece, token_id in vocab.items()}
+        invalid = [token_id for token_id in self.excluded_token_ids if token_id not in id_to_piece]
+        if invalid:
+            raise ValueError(f"excluded token IDs are outside the tokenizer vocabulary: {invalid}")
+        removed_pieces = {id_to_piece[token_id] for token_id in self.excluded_token_ids}
+        model["vocab"] = {
+            piece: token_id for piece, token_id in vocab.items()
+            if int(token_id) not in self.excluded_token_ids
+        }
+
+        def keep_merge(merge):
+            left, right = merge if isinstance(merge, list) else merge.split(" ", 1)
+            return left not in removed_pieces and right not in removed_pieces \
+                and left + right not in removed_pieces
+
+        model["merges"] = [merge for merge in model.get("merges", []) if keep_merge(merge)]
+        config["added_tokens"] = [
+            token for token in config.get("added_tokens", [])
+            if int(token["id"]) not in self.excluded_token_ids
+        ]
+        backend = BackendTokenizer.from_str(json.dumps(config))
+
+        # Clone the Transformers wrapper so normalizers, pre-tokenizers and
+        # return types remain identical, replacing only its Rust backend.
+        input_tokenizer = self.tokenizer.__class__(
+            tokenizer_object=backend,
+            **getattr(self.tokenizer, "init_kwargs", {}),
+        )
+        return input_tokenizer
+
     def _explain_load_error(self, exc):
         """Print a friendly hint when from_pretrained fails on a Hub repo."""
         msg = str(exc)
@@ -378,7 +440,10 @@ def tokenize(self, data):
         # Encode without special tokens to match the behavior of our other
         # subword tokenizers (tiktoken/sentencepiece) so that text resumes
         # cleanly across chunks.
-        ids = self.tokenizer.encode(data, add_special_tokens=False)
+        ids = self.input_tokenizer.encode(data, add_special_tokens=False)
+        reached = sorted(set(ids).intersection(self.excluded_token_ids))
+        if reached:  # Structural invariant and guard against backend changes.
+            raise RuntimeError(f"excluded token IDs were produced by input tokenization: {reached}")
 
         for token_id in ids:
             self.record_token(token_id)
@@ -390,10 +455,15 @@ def tokenize(self, data):
         # gated models like Gemma "just work" downstream: prepare-time is the
         # only step that has to authenticate against the Hub.
         hf_saved_path = None
+        hf_input_saved_path = None
         try:
             os.makedirs(self.hf_local_dir, exist_ok=True)
             self.tokenizer.save_pretrained(self.hf_local_dir)
             hf_saved_path = self.hf_local_dir
+            if self.excluded_token_ids:
+                os.makedirs(self.hf_input_local_dir, exist_ok=True)
+                self.input_tokenizer.save_pretrained(self.hf_input_local_dir)
+                hf_input_saved_path = self.hf_input_local_dir
         except Exception as exc:  # pragma: no cover - best effort
             print(f"[huggingface] Warning: could not save tokenizer snapshot to "
                   f"{self.hf_local_dir}: {exc}")
@@ -403,6 +473,8 @@ def tokenize(self, data):
             "tokenizer": "huggingface",
             "hf_tokenizer_name": self.hf_tokenizer_name,
             "hf_tokenizer_path": hf_saved_path,
+            "hf_input_tokenizer_path": hf_input_saved_path,
+            "hf_excluded_token_ids": self.excluded_token_ids,
             "hf_use_fast": bool(getattr(self.args, "hf_use_fast", True)),
             "hf_trust_remote_code": bool(getattr(self.args, "hf_trust_remote_code", False)),
             "hf_revision": self.hf_revision,

data/template/prepare.py

@@ -61,6 +61,10 @@ def parse_arguments():
                              "Alternatively run `huggingface-cli login` once, or set the "
                              "HF_TOKEN environment variable. You must also accept the model's "
                              "license at https://huggingface.co/<repo> while logged in.")
+    parser.add_argument("--hf_exclude_token_ids", type=int, nargs="*", default=None,
+                        help="Token IDs that user text must never produce. For fast BPE "
+                             "tokenizers (including Gemma 3), the input-only BPE graph is "
+                             "pruned while the original tokenizer remains available for decoding.")
 
     # Sine wave tokenizer arguments
     parser.add_argument("--sine_period", type=float, default=1.0,

data/template/tests.py

@@ -5,6 +5,7 @@
 import sys
 import pickle
 import json
+from unittest.mock import patch
 import prepare
 from nanogpt_tokenizers import (
     SentencePieceTokenizer,
@@ -227,6 +228,43 @@ def test_huggingface_tokenizer(self):
         if os.path.isdir("hf_tokenizer"):
             shutil.rmtree("hf_tokenizer", ignore_errors=True)
 
+    def test_huggingface_tokenizer_excludes_bpe_id_without_breaking_decode(self):
+        try:
+            from tokenizers import Tokenizer as BackendTokenizer
+            from tokenizers.models import BPE
+            from transformers import PreTrainedTokenizerFast
+        except ImportError:
+            self.skipTest("transformers/tokenizers packages not installed")
+
+        backend = BackendTokenizer(BPE(
+            vocab={"a": 0, "b": 1, "ab": 2},
+            merges=[("a", "b")],
+        ))
+        original = PreTrainedTokenizerFast(tokenizer_object=backend)
+        args = Namespace(
+            hf_tokenizer_name="local-test-bpe",
+            hf_trust_remote_code=False,
+            hf_use_fast=True,
+            hf_exclude_token_ids=[2, 2],
+            meta_output_path="meta.pkl",
+        )
+
+        with patch("transformers.AutoTokenizer.from_pretrained", return_value=original):
+            tokenizer = HuggingFaceTokenizer(args)
+        ids = tokenizer.tokenize("ab")
+
+        self.assertEqual(ids, [0, 1])
+        self.assertNotIn(2, ids)
+        self.assertEqual(tokenizer.detokenize([2]), "ab")
+        with open("meta.pkl", "rb") as f:
+            meta = pickle.load(f)
+        self.assertEqual(meta["hf_excluded_token_ids"], [2])
+        self.assertTrue(os.path.isdir(meta["hf_input_tokenizer_path"]))
+
+        import shutil
+        for path in ("hf_tokenizer", "hf_input_tokenizer"):
+            shutil.rmtree(path, ignore_errors=True)
+
 
     def test_custom_tokenizer(self):
         args = Namespace(tokens_file=self.tokens_file)

huggingface_model/gemma/270M/README.md

@@ -18,6 +18,34 @@ The scripts rely on `AutoModelForCausalLM.from_pretrained(...)`, so you can poin
 `--model_name` at any local checkpoint path created by `train_from_scratch.py` or
 `finetune.py`.
 
+## Making token IDs unreachable from user text
+
+Gemma 3's fast tokenizer uses a BPE graph. Deleting an ID from the encoded
+result is unsafe because it silently deletes the text represented by that
+token. The dataset tokenizer instead supports pruning IDs from an **input-only**
+copy of the BPE graph. It removes each selected vocabulary node and all merges
+that use or create that node, so the same text falls back to smaller pieces.
+The original tokenizer is retained for decoding, meaning the model may still
+generate and correctly decode those IDs; only user-supplied text cannot reach
+them.
+
+Exclude one ID or any space-separated list with `--hf_exclude_token_ids`:
+
+```bash
+python data/template/prepare.py \
+  --method huggingface \
+  --hf_tokenizer_name google/gemma-3-270m \
+  --hf_exclude_token_ids 1234 5678 9012 \
+  --train_input input.txt
+```
+
+Preparation saves both tokenizer snapshots next to `meta.pkl`. Inference uses
+the pruned snapshot for prompts and the complete snapshot for generated-output
+decoding. The implementation validates IDs up front and checks every encoded
+result as a defensive invariant. This option requires the fast BPE tokenizer
+(`--hf_use_fast`, which is the default); unsupported model types fail clearly
+rather than falling back to lossy post-processing.
+
 ## JL-projected LM head evaluation
 
 `jl_head_eval.py` runs a two-stage LM head evaluation:

sample.py

@@ -1313,6 +1313,8 @@ def decode_fn(values):
 
         hf_name = meta.get('hf_tokenizer_name')
         hf_local_path = meta.get('hf_tokenizer_path')
+        hf_input_local_path = meta.get('hf_input_tokenizer_path')
+        excluded_token_ids = set(meta.get('hf_excluded_token_ids', []))
         trust_remote_code = bool(meta.get('hf_trust_remote_code', False))
         use_fast = meta.get('hf_use_fast', True)
         hf_revision = meta.get('hf_resolved_commit') or meta.get('hf_revision')
@@ -1345,9 +1347,23 @@ def decode_fn(values):
                 from_pretrained_kwargs["subfolder"] = hf_subfolder
 
         hf_tok = AutoTokenizer.from_pretrained(load_target, **from_pretrained_kwargs)
+        hf_input_tok = hf_tok
+        if excluded_token_ids:
+            if not hf_input_local_path or not os.path.isdir(hf_input_local_path):
+                raise ValueError(
+                    "meta.pkl excludes HuggingFace token IDs but its input-only "
+                    "tokenizer snapshot is missing"
+                )
+            hf_input_tok = AutoTokenizer.from_pretrained(
+                hf_input_local_path, trust_remote_code=trust_remote_code, use_fast=True
+            )
 
         def encode(s):
-            return hf_tok.encode(s, add_special_tokens=False)
+            ids = hf_input_tok.encode(s, add_special_tokens=False)
+            reached = excluded_token_ids.intersection(ids)
+            if reached:
+                raise RuntimeError(f"excluded token IDs reached by user input: {sorted(reached)}")
+            return ids
 
         def decode(ids):
             return hf_tok.decode(list(ids), skip_special_tokens=False)