From 79d88fe8d76c69fa5636e47644ddb96947aa9d6d Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Mon, 3 Aug 2026 09:06:43 -0700 Subject: [PATCH] feat(tokenizer): exclude HuggingFace input token IDs --- data/template/nanogpt_tokenizers.py | 74 +++++++++++++++++++++++++- data/template/prepare.py | 4 ++ data/template/tests.py | 38 +++++++++++++ huggingface_model/gemma/270M/README.md | 28 ++++++++++ sample.py | 18 ++++++- 5 files changed, 160 insertions(+), 2 deletions(-) diff --git a/data/template/nanogpt_tokenizers.py b/data/template/nanogpt_tokenizers.py index 17bba94d42..29418d0d05 100644 --- a/data/template/nanogpt_tokenizers.py +++ b/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, diff --git a/data/template/prepare.py b/data/template/prepare.py index afdd2932dc..bd4fe4d0f2 100644 --- a/data/template/prepare.py +++ b/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/ 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, diff --git a/data/template/tests.py b/data/template/tests.py index 72a3d3e212..93ae09ed18 100644 --- a/data/template/tests.py +++ b/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) diff --git a/huggingface_model/gemma/270M/README.md b/huggingface_model/gemma/270M/README.md index 3553b58dc8..f941ccde65 100644 --- a/huggingface_model/gemma/270M/README.md +++ b/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: diff --git a/sample.py b/sample.py index 378fc51a2d..6c87c9d8d3 100644 --- a/sample.py +++ b/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)