UniqToken

Script-Aware, Entropy-Guided Multilingual Subword Tokenizer

Python tokenizer research toolkit with Rust acceleration, byte fallback, and raw-text span tracking.

What UniqToken Implements

UniqToken is a research tokenizer implementation with trainable Unigram and BPE vocabularies, optional CEM/SuperBPE vocabulary extension, byte fallback, Unicode-aware pre-tokenization, and exact raw-text offset tracking. Imported compatibility models preserve their existing token IDs; research models create a new vocabulary and therefore require a model trained for those IDs.

Token counts depend on the vocabulary, training corpus, normalization, and pre-tokenization configuration. This README does not claim lower API cost, better linguistic boundaries, or superiority over production tokenizers. Those questions require held-out, budget-matched experiments with downstream language models.

Overview

Architecture & Contributor Roadmap: See ROADMAP.md for the eight-stage architecture roadmap and its historical GitHub issue ledger.

UniqToken provides trainable Unigram and BPE models, post-training CEM/SuperBPE vocabulary extension, preprocessing and offset composition, serialization, compatibility importers, and a bundled native Rust extension. Its research-specific mechanisms include script-aware candidate generation and configurable frequency, character-savings, byte-savings, PMI, and boundary-entropy filters. Their empirical effects remain open questions under the protocol below.

Two Engines, One Core

UniqToken's public API is split into two namespaces that share the tokenizer data model and dispatch supported operations to the bundled native Rust core (crates/uniqtoken_core):

# Accelerate an existing model — IDs never change:
from uniqtoken.compat import from_tiktoken

enc = from_tiktoken("cl100k_base.tiktoken", name="cl100k_base", pattern="cl100k_base")

# Train a new vocabulary — research features:
from uniqtoken import UnigramTrainer, SuperBPE

Implementation Contracts

Benchmark Status

UniqToken currently makes no comparative performance or superiority claim. Earlier Phase 14/15 tables, figures, ANOVA results, Pareto analyses, and the pre-integrity matched-budget ledger were produced by harnesses that did not meet the repository's current data-separation and exact-budget contracts. They are retained unchanged under benchmarks/legacy/ for provenance and are not valid evidence for the current implementation.

The active benchmark code now enforces these rules:

• tokenizer training documents are disjoint from every document used for measurement;
• language-model rows identify model_kind explicitly and Transformer evaluation fails if PyTorch or a viable training sequence is unavailable;
• matched trainable tokenizers must reach the exact requested vocabulary size;
• SuperBPE conditions must learn at least one cross-word merge;
• invalid tiers, budgets, devices, and incomplete conditions abort the run instead of producing a partial matched ledger;
• current JSON ledgers carry schema version 3, a full Git commit hash, working-tree dirty status, and a data-split declaration; the matched-budget ledger also records the experiment version. benchmarks.ledger.load_ledger() validates these fields and can require an expected commit hash.

Cross-script density uses tokens_per_unicode_character: emitted token count divided by the number of raw Unicode code points, including whitespace. This replaces whitespace-based fertility for CJK and mixed-script measurements; it is not a linguistic boundary score. The schema-3 loader rejects ambiguous fertility fields and older schemas rather than interpreting them as current results.

Phase A tokenizer screening and the 18-condition Phase B LM screen are complete. Phase B is exploratory screening evidence only; its official interpretation is frozen in benchmarks/PHASE_B_ANALYSIS_REPORT.md. The Phase C confirmatory protocol was frozen but not executed because the required compute exceeded the available free budget. FLORES-200 devtest remained unopened, and the exploratory 16K byte-matched UT-SuperBPE result is not confirmed. See benchmarks/PHASE_C_STATUS.md.

The active entry points are benchmarks/run_phase_a.py for tokenizer stages and benchmarks/run_phase_b_screen.py for the completed one-seed LM screen. benchmarks/run_phase_c_confirm.py implements the frozen confirmation contract but is not launch authorization. The older generic Phase C path in benchmarks/run_research_experiments.py is not authorized for confirmation. Shared ledgers use schema 3; the generic runner uses research schema 5. Dataset, artifact, source, extension, configuration, and completion provenance are validated fail-closed. See benchmarks/RESEARCH_PROTOCOL.md for the execution contracts.

benchmarks/run_matched_budget_eval.py remains a train/validation diagnostic, not a final research experiment. benchmarks/train_toy_transformer.py provides a small three-way train/validation/test sanity harness. benchmarks/downstream_eval.py and benchmarks/benchmark_suite.py report tokenizer-only held-out measurements; they do not establish downstream model quality.

Throughput results are hardware, build, workload, batch-size, and threading dependent. Cross-tokenizer throughput should be compared using input bytes per second because token counts differ by tokenizer. Tokens per second is suitable for comparing implementations only when they produce the same token stream. No throughput table is presented here until a controlled benchmark is rerun from the current HEAD.

The completed Phase B screen does not establish comparative superiority. A publishable confirmatory comparison still requires execution of the frozen three-seed Phase C protocol, its predeclared analysis, held-out test evaluation, uncertainty estimates, and independent reproduction. PAPER_DRAFT.md is a manuscript draft; the versioned files under benchmarks/ are the authoritative experiment protocols, reports, and status records.

Features

Installation

pip install uniqtoken-core==1.0.0

# The distribution exposes the public Python API and native implementation.
python -c "import uniqtoken, uniqtoken_core; print(uniqtoken.__version__)"

# Source checkout / contributor installation
git clone https://github.com/umran666/UniqToken.git
cd UniqToken
pip install -e .

Optional extras (defined in pyproject.toml):

Quickstart

Try UniqToken directly in your browser with our interactive Google Colab Quickstart Tutorial (zero local setup required).

Train a Unigram tokenizer

from uniqtoken import CustomTokenizer

corpus = [...]  # list of training documents

tok = CustomTokenizer.train_from_corpus(
    corpus,
    target_vocab_size=32_000,
    special_tokens=["<|pad|>", "<|unk|>", "<|bos|>", "<|eos|>"],
    byte_fallback=True,
)

# Encode → decode roundtrip (this ASCII example is unchanged by normalization)
ids = tok.encode_to_ids("fix in 2024 at https://site.com")
text = tok.decode(ids)
assert text == "fix in 2024 at https://site.com"

# Stochastic subword regularization (training-time augmentation)
sampled = tok.sample("hello world", alpha=0.5)

# Exact character-span offsets for every token
for token in tok.encode_with_offsets("fix in 2024"):
    print(f"{token.text!r:>12}  id={token.id:<5}  raw_span={token.raw_span}")

Train a BPE tokenizer

from uniqtoken import BPETrainer

trainer = BPETrainer(target_vocab_size=32_000, byte_fallback=True)
model = trainer.train(chunks=corpus, verbose=True)

tokens = model.encode("tokenization")
token_ids = model.encode_to_ids("tokenization")
text = model.decode(token_ids)

Extend vocabulary with CEM / SuperBPE

from uniqtoken import CrossEntropyMerging

# Standard CEM: greedily add merges that minimize cross-entropy increase
cem = CrossEntropyMerging(max_merges=200, verbose=True)
extended = cem.optimize(tok.model, chunks=corpus)

# SuperBPE mode: only accept merges that cross whitespace boundaries
superbpe = CrossEntropyMerging(max_merges=200, cross_word=True)
superbpe_model = superbpe.optimize(tok.model, chunks=corpus)

Export to HuggingFace and GGUF format

# Export to canonical HuggingFace tokenizer.json and tokenizer_config.json
tok.export_to_huggingface("hf_export/")

# Then load with transformers:
# from transformers import AutoTokenizer
# hf_tok = AutoTokenizer.from_pretrained("hf_export/")

# Export to LLaMA.cpp GGUF v3 binary format
tok.export_to_gguf("model.gguf", model_name="llama")

Use the native HuggingFace PreTrainedTokenizerFast adapter

UniqToken ships a transformers.PreTrainedTokenizerFast adapter covering the repository-tested integration surface: save_pretrained / from_pretrained, padding and truncation strategies, return_tensors ("np" and "pt"), batched encoding, and offset mappings. Compatibility outside the tested Transformers versions and surfaces is not guaranteed.

from uniqtoken import CustomTokenizer, UniqTokenizerFast
from transformers import AutoTokenizer

tok = CustomTokenizer.train_from_corpus(corpus, target_vocab_size=8000, verbose=False)

# Wrap the trained tokenizer as a native HF fast tokenizer.
hf_tok = UniqTokenizerFast.from_custom_tokenizer(tok)

# save_pretrained writes tokenizer.json + tokenizer_config.json (with the
# auto_map entry that points back at UniqTokenizerFast), so the repo round-trips
# through the standard HF loaders on any machine with uniqtoken installed.
hf_tok.save_pretrained("uniqtok_export/")

# Reload directly, or let AutoTokenizer discover the custom class.
reloaded = UniqTokenizerFast.from_pretrained("uniqtok_export/")
auto = AutoTokenizer.from_pretrained("uniqtok_export/")  # -> UniqTokenizerFast

Importing uniqtoken.hf_adapter (or accessing uniqtoken.UniqTokenizerFast) registers the class with transformers.AutoTokenizer automatically.

Streaming decode

decoder = tok.get_streaming_decoder()

output = ""
for token_id in generated_ids:  # one id at a time from an LLM
    output += decoder.feed_token_id(token_id)
output += decoder.flush()

Sanitize untrusted input

from uniqtoken import SecurityShield

shield = SecurityShield(special_tokens=["<|endoftext|>", "<|system|>", "<|user|>"])
safe = shield.sanitize(
    untrusted_input,
    allowed_special="none",  # or {"<|user|>"} to whitelist
    disallowed_special_action="escape",  # "escape" | "raise" | "ignore"
)

Note: CustomTokenizer wires SecurityShield.sanitize() into every encode(), sample(), and encode_with_offsets() call automatically (defaults: allowed_special="none", disallowed_special_action="escape"), so sanitization is not an opt-in step.

Compress structured whitespace

from uniqtoken import IndentationCompressor

compact = IndentationCompressor.compress_indents(source_code)
restored = IndentationCompressor.decompress_indents(compact)
assert restored == source_code

Save and load

from uniqtoken import CustomTokenizer

tok.save("saved_model/")
tok2 = CustomTokenizer.load("saved_model/")

assert tok2.encode_to_ids("test") == tok.encode_to_ids("test")

Command-Line Interface (CLI)

UniqToken ships with a CLI executable (uniqtoken) for training, encoding, decoding, and evaluation:

# 1. Train a tokenizer with PMI ranking and SuperBPE optimization
uniqtoken train --corpus dataset.txt --vocab-size 8000 --ranking-strategy pmi --superbpe-merges 100 --out ./model

# 2. Tokenize text with exact character spans and compression telemetry
uniqtoken encode --model ./model --input "def forward(x): return self.attn(x)" --with-metrics

# 3. Encode to integer IDs as JSON
uniqtoken encode --model ./model --input "the quick brown fox" --to-ids --json

# 4. Decode integer IDs to normalized text (NFKC is not raw-byte lossless)
uniqtoken decode --model ./model --input "[12, 450, 89, 230]"

# 5. Run the empirical multilingual benchmark suite with Markdown/LaTeX export
uniqtoken benchmark --export-markdown benchmark_report.md --export-latex table.tex

# 6. Evaluate tokenizer-only context-density proxies on held-out text
uniqtoken eval-downstream --vocab-size 1000

Architecture

End-to-End Pipeline

flowchart LR
    A["Raw Text"] --> B["SecurityShield<br/>sanitize + alignment"]
    B --> C["Normalizer<br/>NFKC + dual-offset"]
    C --> D["RegexPreTokenizer<br/>ordered boundary patterns"]
    D --> E1["UnigramLattice<br/>DAG · Viterbi · FFBS"]
    D --> E2["BPEModel<br/>rank-based merges"]
    E1 --> F["CEM / SuperBPE<br/>vocabulary extension"]
    E1 --> G["Token IDs"]
    E2 --> G
    F --> G
    G --> H["BatchCollator<br/>pad · mask · BOS/EOS"]
    G --> I["StreamingDecoder<br/>byte-buffer aware"]
    H --> J["PyTorch Tensors"]
    I --> K["Decoded Text"]

Project Structure

UniqToken/
├── uniqtoken/                     # Core Python package
│   ├── __init__.py                # Public package namespace & lazy exports
│   ├── cli.py                     # Unified production CLI interface
│   ├── tokenizer.py               # CustomTokenizer — unified facade + parallel batching
│   ├── pre_tokenizer.py           # Normalizer + ordered RegexPreTokenizer boundaries
│   ├── byte_codec.py              # ByteFallbackEngine — UTF-8 ↔ <0xHH> codec
│   ├── trie.py                    # PrefixTrie — prefix lookup for lattice edge mining
│   ├── seed_builder.py            # SeedVocabularyBuilder — PMI + script balancing + entropy
│   ├── unigram_lattice.py         # UnigramLattice — DAG, beam pruning, EM stats, FFBS
│   ├── unigram_trainer.py         # UnigramTrainer — EM early-stopping + Viterbi memoization
│   ├── vocab_adapter.py           # VocabularyAdapter — non-destructive vocab expansion
│   ├── cem_merger.py              # CrossEntropyMerging — CEM / SuperBPE extension
│   ├── bpe_trainer.py             # BPETrainer — classic greedy pairwise-merge training
│   ├── bpe_model.py               # BPEModel — rank-based merge inference (tiktoken-style)
│   ├── batch_collator.py          # BatchCollator — padding, masks, BOS/EOS, to_torch()
│   ├── streaming_decoder.py       # StreamingDecoder — incremental UTF-8-safe decode
│   ├── streaming_counter.py       # Disk-backed counter for bounded-memory training
│   ├── binary_format.py           # Memory-mapped binary model serialization
│   ├── chat_template.py           # Built-in and custom Jinja2 chat templates
│   ├── hf_adapter.py              # Native PreTrainedTokenizerFast adapter
│   ├── hf_exporter.py             # HuggingFaceExporter & GGUFExporter — HF JSON + GGUF v3
│   ├── hf_importer.py             # HuggingFace tokenizer.json importer (Unigram + ByteLevel BPE)
│   ├── sentencepiece_importer.py  # Dependency-free SentencePiece .model protobuf importer
│   ├── tiktoken_adapter.py        # TiktokenEncoding — ranks file loader & exact-ID parity
│   ├── security_shield.py         # SecurityShield — control-token injection defense
│   ├── indentation_compressor.py  # IndentationCompressor — reversible whitespace codec
│   ├── uniqtoken_core.pyi         # Static typing stub for PyO3 native extension
│   ├── integrations/vllm.py       # vLLM-compatible sync/async adapter
│   └── multimodal/                # Multimodal tokenization package
│       ├── __init__.py
│       ├── multimodal_tokenizer.py  # MultimodalTokenizer — text + image
│       ├── visual_codebook.py       # VisualCodebook — VQ codebook for image patches
│       ├── image_patcher.py         # DynamicImagePatcher — grid-based patch extraction
│       ├── audio_codec.py           # Experimental untrained RVQ utility (not supported API)
│       └── neural_codecs.py         # Experimental neural codec building blocks (PyTorch)
│
├── crates/
│   └── uniqtoken_core/            # Native Rust acceleration crate (PyO3 C-extension)
│       ├── Cargo.toml             # Rust package manifest (pyo3, rayon, ahash, regex)
│       └── src/
│           ├── lib.rs             # PyO3 module interface
│           ├── trie.rs            # Native PrefixTrie with AHashMap & prefix search
│           ├── viterbi.rs         # Dynamic programming Viterbi & EM expectations
│           ├── normalizer.rs      # Native Unicode normalization & space handling
│           ├── pipeline.rs        # Fused Rayon batch encoding pipeline
│           ├── rust_tokenizer.rs  # Standalone RustTokenizer engine
│           └── seed.rs            # Native n-gram mining & candidate generation
│
├── benchmarks/
│   ├── benchmark_suite.py                 # Held-out tokenizer compression measurements
│   ├── benchmark_throughput.py            # Byte-normalized implementation throughput
│   ├── downstream_eval.py                 # Held-out tokenizer context-density metrics
│   ├── train_toy_transformer.py           # Small held-out Transformer harness
│   ├── vocab_quality_race.py              # Exact-budget tokenizer comparison harness
│   ├── run_matched_budget_eval.py         # Train/validation diagnostic
│   ├── run_phase_a.py                      # Resumable Phase A screening/selection harness
│   ├── run_phase_b_screen.py               # Frozen-selection Phase B LM screen
│   ├── run_phase_c_confirm.py              # Frozen Phase C contract (not executed)
│   ├── run_research_experiments.py        # Generic accounting/legacy staged interface
│   ├── flop_counter.py                    # Analytical FLOP calculation utilities
│   ├── RESEARCH_PROTOCOL.md                # Authoritative execution contracts
│   ├── PHASE_B_ANALYSIS_REPORT.md          # Frozen exploratory-screen analysis
│   ├── PHASE_C_STATUS.md                   # NOT EXECUTED — COMPUTE CONSTRAINED
│   └── legacy/                            # Invalidated historical scripts, ledgers, figures
│
├── tests/
│   ├── test_tokenizer.py              # Core tokenizer/model behavior
│   ├── test_*compat*.py               # External-format differential compatibility
│   ├── test_phase_*.py                # Phase A/B/C provenance and fail-closed gates
│   ├── test_research_experiments.py   # Accounting and research-schema contracts
│   ├── test_native_*.py               # Rust/Python native pipeline parity
│   └── fuzz/                           # Property-based tokenizer invariants
│
├── assets/banner.jpeg             # Project banner asset
├── CONTRIBUTING.md                # Developer setup and contribution guidelines
├── PAPER_DRAFT.md                 # Research manuscript draft
├── pyproject.toml                 # Package metadata, CLI console_scripts, extras
└── .github/workflows/ci.yml       # CI: 3 OS × 3 Python versions = 9-cell matrix

Module Dependency Graph

graph TD
    CLI["uniqtoken.cli<br/>CLI Commands"] --> T["uniqtoken.tokenizer<br/>CustomTokenizer"]
    T --> N["uniqtoken.pre_tokenizer<br/>Normalizer · RegexPreTokenizer"]
    T --> UL["uniqtoken.unigram_lattice<br/>UnigramLattice"]
    T --> UT["uniqtoken.unigram_trainer<br/>UnigramTrainer · UnigramModel"]
    T --> SS["uniqtoken.security_shield<br/>SecurityShield"]
    T --> IC["uniqtoken.indentation_compressor<br/>IndentationCompressor"]
    T --> SD["uniqtoken.streaming_decoder<br/>StreamingDecoder"]
    T --> HF["uniqtoken.hf_exporter<br/>HuggingFaceExporter · GGUFExporter"]

    UT --> UL
    UT --> SB["uniqtoken.seed_builder<br/>SeedVocabularyBuilder"]
    UT --> BC["uniqtoken.byte_codec<br/>ByteFallbackEngine"]
    UT --> TR["uniqtoken.trie<br/>PrefixTrie"]
    UL --> BC
    UL --> TR
    TR -.-> RC["crates/uniqtoken_core<br/>Rust Native Extension"]
    UL -.-> RC
    T -.-> RC

    CEM["uniqtoken.cem_merger<br/>CrossEntropyMerging"] --> UT
    VA["uniqtoken.vocab_adapter<br/>VocabularyAdapter"] --> UT

    BT["uniqtoken.bpe_trainer<br/>BPETrainer"] --> BC
    BT --> N
    BM["uniqtoken.bpe_model<br/>BPEModel"] --> BC

    MM["uniqtoken.multimodal<br/>MultimodalTokenizer"] --> T

Algorithms & Base Papers

UniqToken is an independent, from-scratch implementation. It does not wrap any paper's reference code. The algorithms are drawn from:

Security Model

SecurityShield guards against control-token smuggling and delimiter hijacking — e.g., a user injecting a literal <|endoftext|> or <|system|> string to manipulate a model's context boundary.

The allowed_special parameter accepts "all", "none", or a specific set of control tokens to whitelist. Sanitization preserves character-alignment tracking via sanitize_with_alignment().

CustomTokenizer integrates this automatically — every encode(), sample(), and encode_with_offsets() call runs through SecurityShield.sanitize() first.

External-Format Compatibility

tiktoken ranks importer

UniqToken loads any tiktoken .tiktoken rank file (e.g. cl100k_base.tiktoken, o200k_base.tiktoken, gpt2 via tiktoken's file dump) and produces exactly the same integer IDs as tiktoken — no tiktoken package required, only the lightweight regex module for pattern fidelity:

from uniqtoken import TiktokenEncoding

enc = TiktokenEncoding.from_file(
    "cl100k_base.tiktoken",
    pattern="cl100k_base",
    special_tokens={"<|endoftext|>": 100257, "<|fim_prefix|>": 100258},
)
ids = enc.encode("Hello, world!")  # identical to tiktoken.encode()
text = enc.decode(ids)

to_uniqtoken_bpe_model() additionally converts the ranks into UniqToken's native BPEModel (IDs preserved) for reuse in training/analysis. CI runs token-for-token differential tests against the real tiktoken package on multilingual, emoji/ZWJ, and code inputs.

HuggingFace tokenizer.json importer

import_hf_tokenizer() reads an HF tokenizer.json (path, directory, or parsed dict) and dispatches on model type:

• Unigram → a native UniqToken CustomTokenizer with scores and token IDs preserved exactly (normalizer/pre-tokenizer mapped best-effort with explicit warnings for unrepresentable components).
• BPE → GPT-2-style ByteLevel vocabs return a fully functional HFByteLevelBPE with exact-ID encode/decode (verified differentially against the real tokenizers package); non-byte-level BPE returns vocab/merges/IDs as a BPEModel for data reuse.
• WordPiece is rejected with a clear error (UniqToken has no WordPiece engine).

from uniqtoken import import_hf_tokenizer

cal = import_hf_tokenizer("path/to/tokenizer.json")  # Unigram -> CustomTokenizer
gpt2 = import_hf_tokenizer("gpt2/tokenizer.json")  # BPE -> HFByteLevelBPE
ids = gpt2.encode("Hello, world!")  # same IDs as HF

Loading a SentencePiece .model (Unigram)

UniqToken can read SentencePiece Unigram models with zero protobuf dependency (raw wire-format parser) and byte-for-byte vocab/ID preservation vs the real sentencepiece package. The first word of every encode is subject to a known SPM/UniqToken divergence (SPM's add_dummy_prefix=True prepends a metaspace that UniqToken does not); the importer emits a UserWarning for it, and the rest of the encode is byte-for-byte identical:

from uniqtoken import import_sentencepiece

tok = import_sentencepiece("sp.model")  # Unigram -> CustomTokenizer
ids = tok.encode_to_ids("hello world")  # IDs preserved; leading-word may differ

Testing & CI

Test Suite

The suite covers core Unigram/BPE behavior, byte fallback, Unicode and offset invariants, native Rust/Python parity, external-format compatibility, CLI/package behavior, property fuzzing, and the fail-closed Phase A/B/C research contracts. Test counts are intentionally not hard-coded here because parametrization and regression coverage change frequently. Use pytest --collect-only -q for the current collected count and pytest for the current pass/skip result.

CI Pipeline

The GitHub Actions workflow runs on every push and PR across a 9-cell matrix (3 OS × 3 Python versions):

Each cell runs:

1. Ruff lint + format check
2. Mypy static type checking
3. Full test suite (unit, adversarial stress, CLI, property fuzzing)
4. Benchmark suite smoke test
5. Package build verification (python -m build)

Running locally

pip install -e ".[test]"

pytest                                          # full test suite
ruff check . && ruff format --check .           # lint + format
mypy uniqtoken                                  # type check
coverage run -m pytest && coverage report       # coverage
python benchmarks/benchmark_suite.py            # benchmark suite
python benchmarks/downstream_eval.py            # held-out tokenizer-only measurements

Multimodal

The multimodal/ package provides experimental text-and-image composition through MultimodalTokenizer. Image token IDs are meaningful only after its visual codebook is trained or a trained codebook is loaded; no trained visual or neural codec checkpoint is bundled. Audio tokenization is unsupported because this distribution does not include a trained audio codebook.

Contributing

1. Fork the repository and create a feature branch.
2. Install the dev toolchain:

   pip install -e ".[test]"
   

3. Keep new code within the ruff (line-length 120, target py310) and mypy configuration.
4. Add or update tests in test_tokenizer.py / test_fuzz_properties.py for any behavioral change.
5. Verify before opening a PR:

   pytest && ruff check . && mypy uniqtoken
   

License

Released under the MIT License.

Maintained by @umran666