tokenizer_inspector.py

#!/usr/bin/env python3
"""
tokenizer_inspector.py v3 — Inspect and compare Hugging Face tokenizers.

Accepts repo IDs (e.g. gpt2, HuggingFaceTB/SmolLM2-135M) or local paths to
tokenizer.json files. Uses the `tokenizers` library for accurate tokenization.

Usage:
    python3 tokenizer_inspector.py gpt2
    python3 tokenizer_inspector.py gpt2 HuggingFaceTB/SmolLM2-135M --compare
    python3 tokenizer_inspector.py gpt2 --text "hello world"
    python3 tokenizer_inspector.py gpt2 --show-tokens

What it reports:
  - Vocab size, token type (BPE / SentencePiece / byte-level)
  - Special tokens
  - chars/token on a test corpus (higher = more compact)
  - Vocab utilization: what fraction of the vocab does your data actually use?
  - Sample of interesting tokens (longest, shortest, most common)
  - In --compare mode: side-by-side comparison table
"""

import argparse
import json
import sys
from collections import Counter
from pathlib import Path

# Default test corpus: mixed English + code + numbers
DEFAULT_CORPUS = """\
The quick brown fox jumps over the lazy dog.
def fibonacci(n):
    if n <= 1: return n
    return fibonacci(n-1) + fibonacci(n-2)
print(fibonacci(10))
#!/usr/bin/env python3
import numpy as np
arr = np.zeros((1024, 1024), dtype=np.float32)
print(arr.sum())
The quick brown fox jumps over the lazy dog.
"""


def load_tokenizer(source: str):
    """Load a tokenizer from a HF repo ID or local path. Returns (Tokenizer, display_name, json_path_or_None)."""
    from tokenizers import Tokenizer

    # Try local path first
    local = Path(source)
    if local.exists() and local.is_file():
        return Tokenizer.from_file(str(local)), local.name, str(local)

    if local.exists() and local.is_dir():
        for name in ("tokenizer.json", "tokenizer.model"):
            p = local / name
            if p.exists():
                return Tokenizer.from_file(str(p)), source, str(p)

    # Try Hugging Face
    try:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(repo_id=source, filename="tokenizer.json")
        return Tokenizer.from_file(path), source, path
    except Exception:
        pass
    try:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(repo_id=source, filename="tokenizer.model")
        return Tokenizer.from_file(path), source, None
    except Exception as e:
        print(f"ERROR: cannot load tokenizer for '{source}': {e}", file=sys.stderr)
        sys.exit(1)


def get_special_tokens(json_path):
    """Read special/added tokens from the tokenizer JSON."""
    if not json_path:
        return []
    try:
        with open(json_path) as f:
            data = json.load(f)
        added = data.get("added_tokens", [])
        return [t["content"] for t in added]
    except Exception:
        return []


def detect_token_type(tok) -> str:
    """Detect the tokenization scheme."""
    s = tok.to_str()
    head = s[:3000]
    if "ByteLevel" in head:
        return "byte-level BPE"
    if "BPE" in head:
        return "BPE"
    if "Unigram" in head:
        return "Unigram (SentencePiece)"
    if "WordPiece" in head:
        return "WordPiece"
    if "WordLevel" in head:
        return "WordLevel"
    return "unknown"


def analyze(source: str, corpus: str, show_tokens: bool = False) -> dict:
    """Analyze one tokenizer and return a metrics dict."""
    tok, name, json_path = load_tokenizer(source)
    vocab_size = tok.get_vocab_size()
    token_type = detect_token_type(tok)
    special_tokens = get_special_tokens(json_path)

    enc = tok.encode(corpus)
    tokens = enc.tokens
    n_tokens = len(tokens)
    n_chars = len(corpus)
    chars_per_token = n_chars / max(n_tokens, 1)

    freq = Counter(tokens)
    most_common = freq.most_common(10)

    # Vocab utilization: how many distinct vocab tokens appear in the corpus?
    distinct_used = set(tokens)
    utilization = len(distinct_used) / max(vocab_size, 1)

    # Longest / shortest real tokens (exclude special)
    vocab = tok.get_vocab()
    real = [t for t in vocab if t and not (t.startswith("<") and t.endswith(">"))]
    longest = max(real, key=len) if real else ("", 0)
    shortest = min((t for t in real if t), key=len, default="")

    # Roundtrip check
    decoded = tok.decode(enc.ids)
    roundtrip_ok = decoded.strip() == corpus.strip()

    result = {
        "repo": name,
        "vocab_size": vocab_size,
        "token_type": token_type,
        "n_special": len(special_tokens),
        "special_sample": special_tokens[:5],
        "corpus_chars": n_chars,
        "corpus_tokens": n_tokens,
        "chars_per_token": round(chars_per_token, 2),
        "vocab_utilization": round(utilization, 4),
        "distinct_tokens_used": len(distinct_used),
        "most_common": most_common,
        "longest_token": (longest[:40], len(longest)),
        "shortest_token": (shortest[:40], len(shortest)),
        "roundtrip_ok": roundtrip_ok,
    }
    if show_tokens:
        result["sample_tokens"] = tokens[:30]
    return result


def print_report(r: dict, show_tokens: bool = False):
    print(f"\n{'='*60}")
    print(f"  {r['repo']}")
    print(f"{'='*60}")
    print(f"  Vocab size:        {r['vocab_size']:,}")
    print(f"  Token type:        {r['token_type']}")
    print(f"  Special tokens:    {r['n_special']} ({', '.join(r['special_sample'])}{'...' if r['n_special'] > 5 else ''})")
    print(f"  Corpus:            {r['corpus_chars']} chars -> {r['corpus_tokens']} tokens")
    print(f"  Chars/token:       {r['chars_per_token']}")
    print(f"  Vocab utilization: {r['vocab_utilization']:.2%} ({r['distinct_tokens_used']} distinct)")
    print(f"  Longest token:     {r['longest_token'][0]!r} ({r['longest_token'][1]} chars)")
    print(f"  Shortest token:    {r['shortest_token'][0]!r} ({r['shortest_token'][1]} chars)")
    print(f"  Roundtrip OK:      {'yes' if r['roundtrip_ok'] else 'NO'}")
    print(f"\n  Top 10 tokens by frequency:")
    for tok, cnt in r["most_common"]:
        print(f"    {repr(tok):>20s}  x{cnt}")
    if show_tokens and "sample_tokens" in r:
        print(f"\n  First 30 tokens:")
        print(f"    {r['sample_tokens']}")
    print()


def print_comparison(results: list):
    print(f"\n{'='*60}")
    print(f"  COMPARISON ({len(results)} tokenizers)")
    print(f"{'='*60}")
    names = [r["repo"][:18] for r in results]
    print(f"  {'Metric':<20s} " + " ".join(f"{n:>20s}" for n in names))
    print(f"  {'-'*20} " + " ".join(f"{'-'*20}" for _ in names))
    rows = [
        ("Vocab size", lambda r: f"{r['vocab_size']:,}"),
        ("Token type", lambda r: r["token_type"]),
        ("Special tokens", lambda r: str(r["n_special"])),
        ("Chars/token", lambda r: f"{r['chars_per_token']}"),
        ("Corpus tokens", lambda r: str(r["corpus_tokens"])),
        ("Vocab util", lambda r: f"{r['vocab_utilization']:.2%}"),
        ("Longest tok", lambda r: str(r["longest_token"][1])),
        ("Roundtrip", lambda r: "ok" if r["roundtrip_ok"] else "FAIL"),
    ]
    for label, fn in rows:
        print(f"  {label:<20s} " + " ".join(f"{fn(r):>20s}" for r in results))
    print()


def main():
    parser = argparse.ArgumentParser(description="Inspect and compare HF tokenizers")
    parser.add_argument("sources", nargs="+", help="Repo IDs or paths to tokenizer.json")
    parser.add_argument("--corpus", type=str, default=DEFAULT_CORPUS, help="Test corpus string")
    parser.add_argument("--text", type=str, default=None, help="Short text to tokenize (overrides default corpus)")
    parser.add_argument("--compare", action="store_true", help="Side-by-side comparison mode")
    parser.add_argument("--show-tokens", action="store_true", help="Show first 30 tokens of the corpus")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    args = parser.parse_args()

    corpus = args.text if args.text else args.corpus
    results = [analyze(s, corpus, args.show_tokens) for s in args.sources]

    if args.json:
        print(json.dumps(results, indent=2))
    elif args.compare and len(results) >= 2:
        print_comparison(results)
        for r in results:
            print_report(r, args.show_tokens)
    else:
        for r in results:
            print_report(r, args.show_tokens)


if __name__ == "__main__":
    main()