Add tokenizer_inspector.py (v3)

tokenizer_inspector.py

@@ -0,0 +1,225 @@
+#!/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()