Unbundle tokenizer.json from the wheel — load from converted-checkpoint dir

[Andrewxu313]

Dispatched to @lausannel · target branch lausannel/unbundle-tokenizer-json-from-the-wheel-l · due 2026-04-25
Push commits here and mark the PR Ready for review once the acceptance criteria are met.

---

Unbundle tokenizer.json from the wheel — load from converted-checkpoint dir

Goal

Stop shipping per-model tokenizer.json files inside the BatchGen wheel. Instead, have the checkpoint converter copy tokenizer.json (and tokenizer_config.json if present) from the source Hugging-Face checkpoint into the converted-checkpoint output directory at conversion time, and make the runtime tokenizer abstraction load those files from the converted-checkpoint directory instead of the bundled package path.

Background

BatchGen currently bundles per-model tokenizer assets as package data:

Model	Path	Size
GPT-OSS-120B	batchgen/models/openai/gpt_oss_120b/tokenizer.json	27 MB
GLM-5	batchgen/models/glm/glm5/tokenizer.json	20 MB
MiniMax-M25	batchgen/models/minimax/minimax_m25/tokenizer.json	9.3 MB
DeepSeek-V3	batchgen/models/deepseek/deepseekv3/tokenizer.json	7.5 MB

Plus four tokenizer_config.json files (~8 KB each). Total: ~64 MB — roughly 14% of the wheel size.

They end up in the wheel because of three cooperating declarations:

• setup.py:160 — package_data={"batchgen": ["**/*.json"]}
• MANIFEST.in:4 — recursive-include batchgen *.json
• pyproject.toml:176 — include_package_data=True

At runtime, each model's tokenizer class resolves its path via Path(__file__).parent and loads tokenizer.json from the package directory. See e.g.:

• batchgen/models/minimax/minimax_m25/tokenizer.py:67 — super().__init__(str(TOKENIZER_DIR))
• batchgen/models/openai/gpt_oss_120b/tokenizer.py:50 — TOKENIZER_DIR = Path(__file__).parent
• batchgen/config/fast_tokenizer.py:73 — tokenizer_file = self.tokenizer_path / "tokenizer.json"

The converted-checkpoint path is already a first-class concept in the server — PR #94 added --converted-ckpt-dir. The server/worker already knows where the converted checkpoint lives at tokenizer-load time.

Design

Two coordinated changes:

1. Converter copies tokenizer.json and tokenizer_config.json from the source checkpoint dir to the converted-checkpoint output dir.
2. Tokenizer loading takes the converted-checkpoint dir as an argument (instead of computing a package-relative path) and reads the tokenizer files from there.

Converter changes

File: batchgen/ckpt_converter/ckpt_converter.py

• In convert(ckpt_path, output_dir, marlin=False) (currently around line 114), after the weights are written, locate tokenizer.json and tokenizer_config.json in the source directory (ckpt_path). If found, copy them verbatim into output_dir.
• If tokenizer.json is missing from the source, print a clear warning naming the missing file and the source path, but do not fail the conversion — some users may convert weights separately from tokenizer setup.
• Do not rename, re-encode, or otherwise modify the files. Byte-identical copy.

File: batchgen/tools/convert_checkpoint.py

• No behavioral change expected; just verify the CLI still works end-to-end after the converter change.

Tokenizer abstraction changes

File: batchgen/config/fast_tokenizer.py

• FastTokenizer.__init__(tokenizer_path) already takes a path — good. Confirm it still works unchanged.

File: batchgen/config/tokenizer_registry.py

• Change load_tokenizer(model_identifier) → load_tokenizer(model_identifier, converted_ckpt_dir). The new argument is the directory containing the converted checkpoint (i.e. where we just copied tokenizer.json).
• Instantiate the resolved model-tokenizer class with the converted-checkpoint path.

Files: each batchgen/models/**/tokenizer.py that currently hard-codes TOKENIZER_DIR = Path(__file__).parent:

• batchgen/models/openai/gpt_oss_120b/tokenizer.py
• batchgen/models/glm/glm5/tokenizer.py
• batchgen/models/deepseek/deepseekv3/tokenizer.py
• batchgen/models/minimax/minimax_m25/tokenizer.py
• (plus any others you find under batchgen/models/*/tokenizer.py)

Change each constructor to accept the tokenizer directory as a parameter and forward it to FastTokenizer.__init__. Remove the TOKENIZER_DIR module-level constant.

File: batchgen/batchgen_worker.py

• Line 1801: self.tokenizer = load_tokenizer(self.huggingface_ckpt_name) → pass the converted-checkpoint directory as well. The worker already knows this path from the --converted-ckpt-dir plumbing added in PR fix: add --converted-ckpt-dir server flag and fix JIT linker paths for conda environments #94; reuse that attribute, don't re-derive it.

File: batchgen/server/batch_scheduler.py

• Line 773: tokenizer = load_tokenizer(model) — same update. Pass the converted-checkpoint dir from the existing server/config plumbing.

Packaging changes

File: setup.py

• Narrow the package_data JSON glob so it no longer captures tokenizer.json / tokenizer_config.json anywhere under batchgen/models/. Keep any other *.json that are genuinely package data (config schemas, model metadata BatchGen itself authors).

File: MANIFEST.in

• Update the recursive-include ... *.json line in tandem with setup.py. Add exclusions for the tokenizer files under batchgen/models/.

File: pyproject.toml

• Review line 176 (include_package_data=True). No change expected if setup.py / MANIFEST.in exclusions are correct.

Delete the bundled tokenizer files from the repo

Remove from the tree (a git rm, not an ignore):

• batchgen/models/openai/gpt_oss_120b/tokenizer.json
• batchgen/models/glm/glm5/tokenizer.json
• batchgen/models/deepseek/deepseekv3/tokenizer.json
• batchgen/models/minimax/minimax_m25/tokenizer.json
• tokenizer_config.json files in the same dirs

Acceptance criteria

• python -m build produces a wheel at least 50 MB smaller than from origin/main. Include the measurement in the PR description.
• Running python -m batchgen.tools.convert_checkpoint --input <hf-checkpoint-dir> --output <out> on a Hugging-Face checkpoint that contains tokenizer.json results in an <out>/tokenizer.json byte-identical to the source (diff -q passes).
• Running the converter on a checkpoint that does not contain tokenizer.json succeeds (exit 0) and prints a clear warning naming the missing file.
• load_tokenizer(model_identifier, converted_ckpt_dir) returns a working tokenizer for each of the four bundled models (GPT-OSS-120B, GLM-5, DeepSeek-V3, MiniMax-M25) when pointed at a converted-checkpoint dir with the copied tokenizer.json.
• A unit test exists at tests/unit/test_tokenizer_loading.py (new file) that, for at least one model, creates a temp dir, copies a small fixture tokenizer.json, calls load_tokenizer, and verifies encode("hello world") returns a non-empty list.
• Existing e2e tests (tests/e2e/gpt_oss_mmlu_pro_test/simple_completion_test.py, tests/e2e/r1_mmlu_pro_test/r1_mmlu_pro_test.py) are updated to point at a converted-checkpoint dir and still pass locally. If you can't run them, say so in the PR and tag the maintainer to verify.
• No Path(__file__).parent-based tokenizer path resolution remains anywhere under batchgen/models/*/tokenizer.py.
• git grep -n tokenizer.json batchgen/ shows no results pointing at a bundled asset (only code references).
• Wheel no longer contains any tokenizer.json — verify with unzip -l dist/*.whl | grep tokenizer.json → empty.

Files expected to change

• batchgen/ckpt_converter/ckpt_converter.py
• batchgen/config/tokenizer_registry.py
• batchgen/models/openai/gpt_oss_120b/tokenizer.py
• batchgen/models/glm/glm5/tokenizer.py
• batchgen/models/deepseek/deepseekv3/tokenizer.py
• batchgen/models/minimax/minimax_m25/tokenizer.py
• (any other batchgen/models/*/tokenizer.py not listed above)
• batchgen/batchgen_worker.py
• batchgen/server/batch_scheduler.py
• setup.py
• MANIFEST.in
• tests/unit/test_tokenizer_loading.py (new)
• tests/e2e/gpt_oss_mmlu_pro_test/simple_completion_test.py (if tokenizer loading touched)
• tests/e2e/r1_mmlu_pro_test/r1_mmlu_pro_test.py (if tokenizer loading touched)
• Deletions: the four tokenizer.json + tokenizer_config.json files under batchgen/models/

How to test

# Build and measure wheel size
python -m build
ls -lh dist/*.whl

# Converter round-trip on a real HF checkpoint that contains tokenizer.json
python -m batchgen.tools.convert_checkpoint \
  --input  /path/to/some-hf-checkpoint \
  --output /tmp/converted-ckpt
test -f /tmp/converted-ckpt/tokenizer.json
diff -q /path/to/some-hf-checkpoint/tokenizer.json /tmp/converted-ckpt/tokenizer.json

# Unit test
pytest tests/unit/test_tokenizer_loading.py -v

# Verify wheel is tokenizer-free
unzip -l dist/*.whl | grep tokenizer.json || echo "OK: no tokenizer.json in wheel"

Out of scope

• Changing the tokenizer encode/decode API. This refactor is about where the files live, not how they're consumed.
• Adding support for tokenizers not currently bundled (new models).
• Marlin INT4 repack behavior (untouched by this task).
• Optimizing wheel size beyond tokenizer.json removal (e.g., other bundled assets).
• Any changes to the --converted-ckpt-dir server flag semantics — reuse it as-is.
• Hot-reload interaction: runtime tokenizer reloading is out of scope; this task is the static path.

Reference PRs

• PR fix: add --converted-ckpt-dir server flag and fix JIT linker paths for conda environments #94 — --converted-ckpt-dir server flag and checkpoint handling. Closest analog to the server-side plumbing you'll be reusing.
• PR feat: batchgen_kernels versioning, ccache, JIT dev mode, and multi-arch support #112 — batchgen_kernels versioning. Example of packaging/setup.py changes.
• c0a6a8b — Marlin repack default. Example of a converter-side modification.

Notes

• The converter's convert() function must be resilient to the case where the user points it at a directory with just .safetensors and no tokenizer.json. Don't hard-fail; warn and continue.
• Keep the copy operation byte-identical. Do not re-serialize the JSON — tokenizers are sensitive to trailing-newline / key-order differences in some cases.
• If you find a model tokenizer file I missed in the list, include it in the same PR and call it out in the PR description.

[lausannel]

• Moved runtime tokenizer assets out of the wheel and into the converted checkpoint directory.
• Updated the checkpoint converter to copy tokenizer assets verbatim from the source checkpoint into converted_ckpt, including support for backfilling missing assets in existing converted directories.
• Changed tokenizer loading so runtime tokenizers read from converted_ckpt instead of package-relative paths.
• Removed bundled tokenizer asset files from the repository and tightened packaging rules so they are no longer included in the wheel.
• Added model-specific tokenizer asset validation in the converter, so required files are checked per model family instead of using one generic rule for all models.
• Updated GPT-OSS to load chat_template.jinja from the converted checkpoint and removed the hardcoded fallback template.
• Updated Kimi K2.5 to load tiktoken.model, tokenizer_config.json, and chat_template.jinja from the converted checkpoint, and replaced the large test fixture with a minimal one.
• Refactored duplicated model-detection logic into a shared lightweight module so model_registry and ckpt_converter use the same detection rules.
• Added a lazy tokenizer import fallback in tokenizer_registry to avoid registration failures caused by heavy package import side effects.
• Expanded unit tests to cover asset copying, runtime tokenizer loading, model-specific validation, GPT-OSS template handling, and Kimi fixtures.