app.py is updated. tokenizer_source_block() function is reused across all 3 tabs (Playground, Evaluate, Compare) and renders a radio toggle with three panels that show/hide dynamically

README.md

@@ -42,4 +42,4 @@ Evaluate any Hugging Face or tiktoken tokenizer against the **TokenizerBench** d
 ## Supported tokenizer types
 
 - **HuggingFace AutoTokenizer** — any model from the Hub, e.g. `bert-base-multilingual-cased`, `xlm-roberta-base`, `google/mt5-base`
-- **tiktoken** — OpenAI encodings: `cl100k_base`, `o200k_base`, `p50k_base`
+- **tiktoken** — OpenAI encodings: `cl100k_base`, `o200k_base`, `p50k_base`

app.py

@@ -1,11 +1,12 @@
 """
 TokenizerBench — Hugging Face Space
-A Gradio app that lets users try any HF/tiktoken tokenizer against
-the TokenizerBench dataset and visualise the results.
+Evaluate and compare tokenizers on the TokenizerBench dataset.
+Supports: HuggingFace AutoTokenizer (Hub ID or uploaded files),
+          tiktoken encodings, SentencePiece .model files.
 """
 
 import io
-import json
+import shutil
 import tempfile
 import traceback
 from pathlib import Path
@@ -21,7 +22,7 @@ import pandas as pd
 matplotlib.use("Agg")
 
 # ─────────────────────────────────────────────────────────────────
-# Inline dataset (subset of the full TokenizerBench data)
+# Dataset (inline subset of TokenizerBench)
 # ─────────────────────────────────────────────────────────────────
 
 DATASET: dict[str, dict[str, list[str]]] = {
@@ -64,8 +65,8 @@ DATASET: dict[str, dict[str, list[str]]] = {
         "german": [
             "Künstliche Intelligenz verändert die Welt schnell.",
             "Ich lerne gerne neue Technologien.",
-            "Dies ist ein Testsatz.",
             "Donaudampfschifffahrtsgesellschaft ist ein langes deutsches Wort.",
+            "Dies ist ein Testsatz.",
             "Natürliche Sprachverarbeitung ist ein wichtiges Forschungsgebiet.",
         ],
         "russian": [
@@ -96,13 +97,11 @@ DATASET: dict[str, dict[str, list[str]]] = {
             "const nums = [1,2,3]; const sq = nums.map(x => x**2);",
             "async function fetchData(url) { const res = await fetch(url); return res.json(); }",
             "const obj = { key: 'value', nested: { a: 1 } };",
-            "document.querySelector('#app').innerHTML = '<h1>Hello</h1>';",
         ],
         "sql": [
             "SELECT u.name, COUNT(o.id) FROM users u JOIN orders o ON u.id = o.user_id GROUP BY u.name;",
             "CREATE INDEX idx_users_email ON users(email);",
             "WITH ranked AS (SELECT *, ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) rn FROM emp) SELECT * FROM ranked WHERE rn=1;",
-            "INSERT INTO logs (event, ts) VALUES ('login', NOW());",
         ],
         "rust": [
             "fn main() { println!(\"Hello, world!\"); }",
@@ -152,7 +151,7 @@ DATASET: dict[str, dict[str, list[str]]] = {
         "mixed_scripts": [
             "Hello 世界 مرحبا Привет こんにちは",
             "AI模型 and NLP技术 are transforming الذكاء الاصطناعي",
-            "math: α + β = γ, code: x += 1, emoji: 🚀",
+            "math: α + β = γ, code: x += 1",
         ],
     },
 }
@@ -165,62 +164,140 @@ CATEGORY_LABELS = {
 }
 
 # ─────────────────────────────────────────────────────────────────
-# Metrics (mirrors metrics.py from the repo)
+# Tokenizer loaders
+# ─────────────────────────────────────────────────────────────────
+
+def _hf_wrapper(tok):
+    class W:
+        def encode(self, text):
+            return tok.encode(text, add_special_tokens=False)
+        def decode(self, ids):
+            return tok.decode(ids, skip_special_tokens=True,
+                              clean_up_tokenization_spaces=False)
+    return W()
+
+
+def load_from_hub(model_id: str):
+    from transformers import AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(model_id.strip())
+    return _hf_wrapper(tok), model_id.strip()
+
+
+def load_from_uploaded_files(files: list, display_name: str):
+    """
+    Accepts a list of Gradio file objects and returns (wrapper, name).
+
+    Supported combinations:
+      • tokenizer.json  [+ tokenizer_config.json, vocab.txt, merges.txt …]
+        → HuggingFace fast tokenizer loaded from a temp dir
+      • *.model
+        → SentencePiece
+      • vocab.json + merges.txt   (BPE without tokenizer.json)
+        → HuggingFace from temp dir
+    """
+    if not files:
+        raise ValueError("No files uploaded.")
+
+    paths = [Path(f.name) for f in files]
+    filenames = {p.name for p in paths}
+
+    # ── SentencePiece .model ───────────────────────────────────
+    sp_models = [p for p in paths if p.suffix == ".model"]
+    if sp_models:
+        import sentencepiece as spm
+        sp = spm.SentencePieceProcessor()
+        sp.Load(str(sp_models[0]))
+        class SPWrapper:
+            def encode(self, text): return sp.EncodeAsIds(text)
+            def decode(self, ids): return sp.DecodeIds(ids)
+        return SPWrapper(), display_name or sp_models[0].stem
+
+    # ── HuggingFace file set ───────────────────────────────────
+    HF_FILES = {
+        "tokenizer.json", "tokenizer_config.json",
+        "vocab.txt", "vocab.json", "merges.txt",
+        "special_tokens_map.json", "added_tokens.json", "spiece.model",
+    }
+    hf_uploads = [p for p in paths if p.name in HF_FILES]
+    if hf_uploads:
+        from transformers import AutoTokenizer
+        tmp = Path(tempfile.mkdtemp(prefix="tok_"))
+        for p in hf_uploads:
+            shutil.copy(p, tmp / p.name)
+        tok = AutoTokenizer.from_pretrained(str(tmp))
+        return _hf_wrapper(tok), display_name or "uploaded-tokenizer"
+
+    raise ValueError(
+        f"Unrecognised file(s): {', '.join(p.name for p in paths)}.\n"
+        "Expected: tokenizer.json, *.model, or vocab.json + merges.txt"
+    )
+
+
+def load_tiktoken(encoding: str):
+    import tiktoken
+    enc = tiktoken.get_encoding(encoding.strip())
+    class W:
+        def encode(self, text): return enc.encode(text)
+        def decode(self, ids): return enc.decode(ids)
+    return W(), encoding.strip()
+
+
+def resolve_tokenizer(source, hub_id, uploaded_files, upload_name, tiktoken_enc):
+    if source == "HuggingFace Hub ID":
+        if not hub_id.strip():
+            raise ValueError("Please enter a Hub model ID (e.g. bert-base-multilingual-cased).")
+        return load_from_hub(hub_id)
+    elif source == "Upload files":
+        if not uploaded_files:
+            raise ValueError("Please upload at least one tokenizer file.")
+        return load_from_uploaded_files(uploaded_files, (upload_name or "").strip())
+    elif source == "tiktoken encoding":
+        if not tiktoken_enc.strip():
+            raise ValueError("Please enter a tiktoken encoding (e.g. cl100k_base).")
+        return load_tiktoken(tiktoken_enc)
+    raise ValueError(f"Unknown source: {source}")
+
+# ─────────────────────────────────────────────────────────────────
+# Metrics
 # ─────────────────────────────────────────────────────────────────
 
-def fertility_score(tokenizer, text: str) -> float:
+def fertility_score(tok, text):
     words = text.split()
-    if not words:
-        return 0.0
-    tokens = tokenizer.encode(text)
-    return len(tokens) / len(words)
-
-def compression_ratio(tokenizer, text: str) -> float:
-    if not text:
-        return 0.0
-    return len(tokenizer.encode(text)) / len(text)
-
-def byte_compression_ratio(tokenizer, text: str) -> float:
-    n_bytes = len(text.encode("utf-8"))
-    if n_bytes == 0:
-        return 0.0
-    return len(tokenizer.encode(text)) / n_bytes
-
-def roundtrip_fidelity(tokenizer, text: str) -> bool:
+    return len(tok.encode(text)) / len(words) if words else 0.0
+
+def compression_ratio(tok, text):
+    return len(tok.encode(text)) / len(text) if text else 0.0
+
+def byte_compression_ratio(tok, text):
+    n = len(text.encode("utf-8"))
+    return len(tok.encode(text)) / n if n else 0.0
+
+def roundtrip_fidelity(tok, text):
     try:
-        ids = tokenizer.encode(text)
-        decoded = tokenizer.decode(ids)
-        return text.strip() == decoded.strip()
+        return text.strip() == tok.decode(tok.encode(text)).strip()
     except Exception:
         return False
 
-def evaluate_tokenizer(tokenizer, dataset: dict) -> dict:
+def evaluate_tokenizer(tok, dataset):
     results: dict[str, Any] = {}
-    all_f, all_c = [], []
-    failures = 0
-
+    all_f, all_c, failures = [], [], 0
     for category, subcategories in dataset.items():
         results[category] = {}
         for subcategory, samples in subcategories.items():
-            ferts, comps, byte_comps, token_counts = [], [], [], []
-            sub_fails = 0
+            ferts, comps, byte_comps, token_counts, sub_fails = [], [], [], [], 0
             for text in samples:
                 if not text or not text.strip():
                     continue
                 try:
-                    toks = tokenizer.encode(text)
-                    token_counts.append(len(toks))
-                    f = fertility_score(tokenizer, text)
-                    ferts.append(f); all_f.append(f)
-                    c = compression_ratio(tokenizer, text)
-                    comps.append(c); all_c.append(c)
-                    byte_comps.append(byte_compression_ratio(tokenizer, text))
-                    if not roundtrip_fidelity(tokenizer, text):
+                    token_counts.append(len(tok.encode(text)))
+                    f = fertility_score(tok, text); ferts.append(f); all_f.append(f)
+                    c = compression_ratio(tok, text); comps.append(c); all_c.append(c)
+                    byte_comps.append(byte_compression_ratio(tok, text))
+                    if not roundtrip_fidelity(tok, text):
                         sub_fails += 1; failures += 1
                 except Exception:
                     pass
-
-            def avg(lst): return round(sum(lst)/len(lst), 4) if lst else 0.0
+            def avg(l): return round(sum(l)/len(l), 4) if l else 0.0
             results[category][subcategory] = {
                 "n_samples": len(token_counts),
                 "avg_tokens": avg(token_counts),
@@ -229,7 +306,6 @@ def evaluate_tokenizer(tokenizer, dataset: dict) -> dict:
                 "avg_byte_compression": avg(byte_comps),
                 "fidelity_failures": sub_fails,
             }
-
     results["__summary__"] = {
         "overall_avg_fertility": round(sum(all_f)/len(all_f), 4) if all_f else 0,
         "overall_avg_compression": round(sum(all_c)/len(all_c), 4) if all_c else 0,
@@ -238,34 +314,12 @@ def evaluate_tokenizer(tokenizer, dataset: dict) -> dict:
     }
     return results
 
-# ─────────────────────────────────────────────────────────────────
-# Tokenizer loaders
-# ─────────────────────────────────────────────────────────────────
-
-def load_hf_tokenizer(model_id: str):
-    from transformers import AutoTokenizer
-    tok = AutoTokenizer.from_pretrained(model_id)
-    class W:
-        def encode(self, text):
-            return tok.encode(text, add_special_tokens=False)
-        def decode(self, ids):
-            return tok.decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
-    return W()
-
-def load_tiktoken(model: str):
-    import tiktoken
-    enc = tiktoken.get_encoding(model)
-    class W:
-        def encode(self, text): return enc.encode(text)
-        def decode(self, ids): return enc.decode(ids)
-    return W()
-
 # ─────────────────────────────────────────────────────────────────
 # Plots
 # ─────────────────────────────────────────────────────────────────
 
-PALETTE = ["#3b82f6", "#8b5cf6", "#ec4899", "#f59e0b", "#10b981",
-           "#ef4444", "#06b6d4", "#84cc16"]
+PALETTE = ["#3b82f6","#8b5cf6","#ec4899","#f59e0b","#10b981",
+           "#ef4444","#06b6d4","#84cc16"]
 
 def fig_to_pil(fig):
     buf = io.BytesIO()
@@ -275,18 +329,20 @@ def fig_to_pil(fig):
     from PIL import Image
     return Image.open(buf).copy()
 
-def plot_fertility_heatmap(result: dict, title: str):
+def _dark_fig(w, h):
+    fig, ax = plt.subplots(figsize=(w, h), facecolor="#0f1117")
+    ax.set_facecolor("#0f1117")
+    ax.tick_params(colors="white")
+    ax.spines[["top","right","bottom","left"]].set_color("#333")
+    return fig, ax
+
+def plot_fertility_heatmap(result, title):
     cats = [c for c in result if not c.startswith("__") and isinstance(result[c], dict)]
-    if not cats:
-        return None
-    data = {}
-    for cat in cats:
-        data[cat] = {sub: v.get("avg_fertility", 0)
-                     for sub, v in result[cat].items() if isinstance(v, dict)}
+    if not cats: return None
+    data = {cat: {sub: v.get("avg_fertility", 0) for sub, v in result[cat].items()
+                  if isinstance(v, dict)} for cat in cats}
     df = pd.DataFrame(data).T.fillna(0)
-    fig, ax = plt.subplots(figsize=(max(10, len(df.columns)*0.8), max(4, len(df)*0.6)),
-                           facecolor="#0f1117")
-    ax.set_facecolor("#0f1117")
+    fig, ax = _dark_fig(max(10, len(df.columns)*0.8), max(4, len(df)*0.6))
     import seaborn as sns
     sns.heatmap(df, ax=ax, cmap="YlOrRd", annot=True, fmt=".2f",
                 linewidths=0.5, linecolor="#1e2130",
@@ -298,55 +354,42 @@ def plot_fertility_heatmap(result: dict, title: str):
     ax.figure.axes[-1].tick_params(colors="white", labelsize=8)
     ax.figure.axes[-1].yaxis.label.set_color("white")
     plt.tight_layout()
-    img = fig_to_pil(fig)
-    plt.close(fig)
-    return img
+    img = fig_to_pil(fig); plt.close(fig); return img
 
-def plot_language_fertility_bar(result: dict, title: str):
+def plot_language_fertility_bar(result, title):
     lang_data = result.get("human_languages", {})
-    if not lang_data:
-        return None
+    if not lang_data: return None
     langs = {lang: v["avg_fertility"] for lang, v in lang_data.items()
              if isinstance(v, dict) and "avg_fertility" in v}
     langs = dict(sorted(langs.items(), key=lambda x: x[1]))
-    colors = ["#d73027" if v > 3 else "#fdae61" if v > 2 else "#1a9850"
-              for v in langs.values()]
-    fig, ax = plt.subplots(figsize=(9, max(4, len(langs)*0.35)), facecolor="#0f1117")
-    ax.set_facecolor("#0f1117")
+    colors = ["#d73027" if v > 3 else "#fdae61" if v > 2 else "#1a9850" for v in langs.values()]
+    fig, ax = _dark_fig(9, max(4, len(langs)*0.35))
     bars = ax.barh(list(langs.keys()), list(langs.values()), color=colors, height=0.7)
     for bar, val in zip(bars, langs.values()):
-        ax.text(val + 0.02, bar.get_y() + bar.get_height()/2,
+        ax.text(val+0.02, bar.get_y()+bar.get_height()/2,
                 f"{val:.2f}", va="center", fontsize=8, color="white")
-    ax.axvline(1.0, color="#aaa", linestyle="--", linewidth=0.8, label="Ideal (1.0)")
-    ax.axvline(2.0, color="#fdae61", linestyle="--", linewidth=0.8, label="Acceptable (2.0)")
-    ax.axvline(4.0, color="#d73027", linestyle="--", linewidth=0.8, label="Poor (≥4.0)")
+    ax.axvline(1.0, color="#aaa", linestyle="--", lw=0.8, label="Ideal (1.0)")
+    ax.axvline(2.0, color="#fdae61", linestyle="--", lw=0.8, label="Acceptable (2.0)")
+    ax.axvline(4.0, color="#d73027", linestyle="--", lw=0.8, label="Poor (≥4.0)")
     ax.set_xlabel("Avg fertility (tokens/word)", color="white")
     ax.set_title(f"Per-language fertility — {title}", color="white", fontsize=11)
-    ax.tick_params(colors="white", labelsize=9)
-    ax.spines[["top","right","bottom","left"]].set_color("#333")
-    legend = ax.legend(fontsize=8, facecolor="#1e2130", labelcolor="white")
+    ax.legend(fontsize=8, facecolor="#1e2130", labelcolor="white")
     plt.tight_layout()
-    img = fig_to_pil(fig)
-    plt.close(fig)
-    return img
+    img = fig_to_pil(fig); plt.close(fig); return img
 
-def plot_compression_scatter(result: dict, title: str):
+def plot_compression_scatter(result, title):
     xs, ys, labels, cat_list = [], [], [], []
     cat_colors = {}
     cats = [c for c in result if not c.startswith("__") and isinstance(result[c], dict)]
     for i, cat in enumerate(cats):
         cat_colors[cat] = PALETTE[i % len(PALETTE)]
         for sub, vals in result[cat].items():
-            if not isinstance(vals, dict):
-                continue
+            if not isinstance(vals, dict): continue
             f = vals.get("avg_fertility"); c = vals.get("avg_byte_compression")
             if f is not None and c is not None:
-                xs.append(c); ys.append(f)
-                labels.append(sub); cat_list.append(cat)
-    if not xs:
-        return None
-    fig, ax = plt.subplots(figsize=(9, 6), facecolor="#0f1117")
-    ax.set_facecolor("#0f1117")
+                xs.append(c); ys.append(f); labels.append(sub); cat_list.append(cat)
+    if not xs: return None
+    fig, ax = _dark_fig(9, 6)
     for cat in set(cat_list):
         idxs = [i for i, c in enumerate(cat_list) if c == cat]
         ax.scatter([xs[i] for i in idxs], [ys[i] for i in idxs],
@@ -355,130 +398,196 @@ def plot_compression_scatter(result: dict, title: str):
     for i, lbl in enumerate(labels):
         ax.annotate(lbl, (xs[i], ys[i]), fontsize=6.5, color="#ccc",
                     xytext=(4, 3), textcoords="offset points")
-    ax.axhline(1.0, color="#aaa", linestyle="--", linewidth=0.8, label="Fertility=1.0")
-    ax.axhline(2.0, color="#fdae61", linestyle="--", linewidth=0.8, label="Fertility=2.0")
+    ax.axhline(1.0, color="#aaa", linestyle="--", lw=0.8, label="Fertility=1.0")
+    ax.axhline(2.0, color="#fdae61", linestyle="--", lw=0.8, label="Fertility=2.0")
     ax.set_xlabel("Byte compression (tokens/byte) — lower is better", color="white")
     ax.set_ylabel("Fertility (tokens/word) — lower is better", color="white")
     ax.set_title(f"Fertility vs byte compression — {title}", color="white", fontsize=11)
-    ax.tick_params(colors="white")
-    ax.spines[["top","right","bottom","left"]].set_color("#333")
     ax.legend(fontsize=8, facecolor="#1e2130", labelcolor="white")
     plt.tight_layout()
-    img = fig_to_pil(fig)
-    plt.close(fig)
-    return img
-
-def plot_comparison_bar(results_dict: dict, metric: str = "avg_fertility"):
-    if not results_dict:
-        return None
-    cats = set()
-    data: dict[str, dict[str, float]] = {}
+    img = fig_to_pil(fig); plt.close(fig); return img
+
+def plot_comparison_bar(results_dict, metric="avg_fertility"):
+    if not results_dict: return None
+    cats, data = set(), {}
     for tok_name, result in results_dict.items():
         data[tok_name] = {}
         for cat, subcats in result.items():
-            if cat.startswith("__") or not isinstance(subcats, dict):
-                continue
+            if cat.startswith("__") or not isinstance(subcats, dict): continue
             vals = [v.get(metric, 0) for v in subcats.values()
                     if isinstance(v, dict) and metric in v]
             if vals:
-                data[tok_name][cat] = round(sum(vals)/len(vals), 4)
-                cats.add(cat)
+                data[tok_name][cat] = round(sum(vals)/len(vals), 4); cats.add(cat)
     cats = sorted(cats)
     tok_names = list(data.keys())
     x = np.arange(len(cats))
     width = 0.75 / max(len(tok_names), 1)
-    fig, ax = plt.subplots(figsize=(max(9, len(cats)*1.8), 5.5), facecolor="#0f1117")
-    ax.set_facecolor("#0f1117")
+    fig, ax = _dark_fig(max(9, len(cats)*1.8), 5.5)
     for i, name in enumerate(tok_names):
         vals = [data[name].get(cat, 0) for cat in cats]
         offset = x + i*width - (len(tok_names)-1)*width/2
         bars = ax.bar(offset, vals, width*0.9, label=name,
                       color=PALETTE[i % len(PALETTE)], alpha=0.88)
         for bar, val in zip(bars, vals):
-            ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
+            ax.text(bar.get_x()+bar.get_width()/2, bar.get_height()+0.01,
                     f"{val:.2f}", ha="center", va="bottom", fontsize=7.5, color="white")
-    cat_labels = [CATEGORY_LABELS.get(c, c) for c in cats]
     ax.set_xticks(x)
-    ax.set_xticklabels(cat_labels, rotation=20, ha="right", color="white", fontsize=9)
-    ax.set_ylabel(metric.replace("_", " ").title(), color="white")
+    ax.set_xticklabels([CATEGORY_LABELS.get(c, c) for c in cats],
+                       rotation=20, ha="right", color="white", fontsize=9)
+    ax.set_ylabel(metric.replace("_"," ").title(), color="white")
     ax.set_title(f"Tokenizer comparison — {metric.replace('_',' ').title()}", color="white", fontsize=11)
-    ax.tick_params(colors="white")
-    ax.spines[["top","right","bottom","left"]].set_color("#333")
     ax.legend(fontsize=9, facecolor="#1e2130", labelcolor="white")
     plt.tight_layout()
-    img = fig_to_pil(fig)
-    plt.close(fig)
-    return img
+    img = fig_to_pil(fig); plt.close(fig); return img
 
-def plot_fidelity_summary(results_dict: dict):
+def plot_fidelity_summary(results_dict):
     names = list(results_dict.keys())
     failures = [r.get("__summary__", {}).get("fidelity_failure_count", 0)
                 for r in results_dict.values()]
-    fig, ax = plt.subplots(figsize=(max(5, len(names)*1.4), 4.5), facecolor="#0f1117")
-    ax.set_facecolor("#0f1117")
+    fig, ax = _dark_fig(max(5, len(names)*1.4), 4.5)
     colors = ["#d73027" if f > 0 else "#1a9850" for f in failures]
     bars = ax.bar(names, failures, color=colors, width=0.5)
     for bar, val in zip(bars, failures):
-        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.05,
+        ax.text(bar.get_x()+bar.get_width()/2, bar.get_height()+0.05,
                 str(val), ha="center", va="bottom", fontsize=10,
                 color="#d73027" if val > 0 else "#1a9850")
     ax.set_ylabel("Fidelity failure count", color="white")
     ax.set_title("Roundtrip fidelity failures", color="white", fontsize=11)
-    ax.tick_params(colors="white")
-    ax.spines[["top","right","bottom","left"]].set_color("#333")
     ax.set_ylim(bottom=0)
-    green_patch = mpatches.Patch(color="#1a9850", label="0 failures (pass)")
-    red_patch = mpatches.Patch(color="#d73027", label="Has failures")
-    ax.legend(handles=[green_patch, red_patch], fontsize=8,
-              facecolor="#1e2130", labelcolor="white")
+    ax.legend(handles=[
+        mpatches.Patch(color="#1a9850", label="0 failures (pass)"),
+        mpatches.Patch(color="#d73027", label="Has failures"),
+    ], fontsize=8, facecolor="#1e2130", labelcolor="white")
     plt.tight_layout()
-    img = fig_to_pil(fig)
-    plt.close(fig)
-    return img
+    img = fig_to_pil(fig); plt.close(fig); return img
 
 # ─────────────────────────────────────────────────────────────────
-# Core Gradio logic
+# Shared tokenizer source block builder
 # ─────────────────────────────────────────────────────────────────
 
-def run_single_eval(model_id: str, tok_type: str, categories: list[str]):
-    if not model_id.strip():
-        return "⚠️ Please enter a model name.", None, None, None, None
+def tokenizer_source_block(prefix=""):
+    """Renders the three-way tokenizer source UI and returns component dict."""
+    gr.Markdown(f"#### {prefix}Load tokenizer")
+
+    source = gr.Radio(
+        ["HuggingFace Hub ID", "Upload files", "tiktoken encoding"],
+        value="HuggingFace Hub ID",
+        label="Source",
+    )
+
+    with gr.Column(visible=True) as hub_col:
+        hub_id = gr.Textbox(
+            label="Hub model ID",
+            placeholder="bert-base-multilingual-cased",
+            value="bert-base-multilingual-cased",
+        )
+        gr.Markdown(
+            "<small>Examples: `xlm-roberta-base` · `google/mt5-base` · "
+            "`facebook/mbart-large-50` · `ai4bharat/indic-bert`</small>"
+        )
+
+    with gr.Column(visible=False) as upload_col:
+        uploaded_files = gr.File(
+            label="Upload tokenizer file(s)",
+            file_count="multiple",
+            file_types=[".json", ".txt", ".model", ".bpe", ".vocab"],
+        )
+        upload_name = gr.Textbox(
+            label="Display name (optional)",
+            placeholder="my-custom-tokenizer",
+        )
+        gr.Markdown(
+            "<small>"
+            "**HuggingFace fast tokenizer** → upload `tokenizer.json` "
+            "(optionally also `tokenizer_config.json`, `vocab.txt`, `merges.txt`)<br>"
+            "**SentencePiece** → upload the `.model` file<br>"
+            "**BPE (GPT-2 style)** → upload `vocab.json` + `merges.txt`"
+            "</small>"
+        )
+
+    with gr.Column(visible=False) as tiktoken_col:
+        tiktoken_enc = gr.Textbox(
+            label="Encoding name",
+            placeholder="cl100k_base",
+            value="cl100k_base",
+        )
+        gr.Markdown(
+            "<small>Available encodings: "
+            "`cl100k_base` (GPT-3.5/4) · `o200k_base` (GPT-4o) · `p50k_base` (Codex)</small>"
+        )
+
+    # dummy defaults so every branch always has a value
+    hub_id_default = gr.Textbox(value="", visible=False)
+    upload_name_default = gr.Textbox(value="", visible=False)
+    tiktoken_enc_default = gr.Textbox(value="cl100k_base", visible=False)
 
-    status = ""
+    def _toggle(s):
+        return (
+            gr.update(visible=s == "HuggingFace Hub ID"),
+            gr.update(visible=s == "Upload files"),
+            gr.update(visible=s == "tiktoken encoding"),
+        )
+    source.change(_toggle, source, [hub_col, upload_col, tiktoken_col])
+
+    return dict(
+        source=source,
+        hub_id=hub_id,
+        uploaded_files=uploaded_files,
+        upload_name=upload_name,
+        tiktoken_enc=tiktoken_enc,
+    )
+
+# ─────────────────────────────────────────────────────────────────
+# Tab logic
+# ─────────────────────────────────────────────────────────────────
+
+def tokenize_live(source, hub_id, uploaded_files, upload_name, tiktoken_enc, text):
+    if not text.strip():
+        return "Enter some text above to tokenize.", ""
     try:
-        if tok_type == "HuggingFace (AutoTokenizer)":
-            tok = load_hf_tokenizer(model_id.strip())
-        else:
-            tok = load_tiktoken(model_id.strip())
-    except Exception as e:
-        return f"❌ Failed to load tokenizer:\n{traceback.format_exc()}", None, None, None, None
-
-    dataset_subset = {k: v for k, v in DATASET.items() if k in categories} if categories else DATASET
-    if not dataset_subset:
-        return "⚠️ Please select at least one dataset category.", None, None, None, None
+        tok, name = resolve_tokenizer(source, hub_id, uploaded_files, upload_name, tiktoken_enc)
+    except Exception:
+        return f"❌ Could not load tokenizer:\n```\n{traceback.format_exc()}\n```", ""
+    try:
+        ids = tok.encode(text)
+        fid = "✅ Roundtrip OK" if roundtrip_fidelity(tok, text) else "⚠️ Roundtrip mismatch"
+        info = (
+            f"**Tokenizer:** `{name}`  \n"
+            f"**Token count:** {len(ids)}  |  "
+            f"**Fertility:** {len(ids)/max(1,len(text.split())):.2f}  |  "
+            f"**Compression:** {len(ids)/max(1,len(text)):.3f}  |  {fid}"
+        )
+        ids_str = " ".join(str(i) for i in ids[:120])
+        if len(ids) > 120:
+            ids_str += f"  … (+{len(ids)-120} more)"
+        return info, ids_str
+    except Exception:
+        return f"❌ Tokenization error:\n```\n{traceback.format_exc()}\n```", ""
+
 
+def run_single_eval(source, hub_id, uploaded_files, upload_name, tiktoken_enc, categories):
+    try:
+        tok, name = resolve_tokenizer(source, hub_id, uploaded_files, upload_name, tiktoken_enc)
+    except Exception:
+        return f"❌ Could not load tokenizer:\n```\n{traceback.format_exc()}\n```", None, None, None, None
+    dataset_subset = {k: v for k, v in DATASET.items() if k in (categories or [])}
+    if not dataset_subset:
+        return "⚠️ Select at least one dataset category.", None, None, None, None
     try:
         result = evaluate_tokenizer(tok, dataset_subset)
-    except Exception as e:
-        return f"❌ Evaluation error:\n{traceback.format_exc()}", None, None, None, None
-
+    except Exception:
+        return f"❌ Evaluation error:\n```\n{traceback.format_exc()}\n```", None, None, None, None
     s = result["__summary__"]
     status = (
-        f"✅ **{model_id.strip()}** evaluated on {s['total_samples']} samples\n\n"
+        f"✅ **{name}** — {s['total_samples']} samples evaluated\n\n"
         f"| Metric | Value |\n|--------|-------|\n"
         f"| Overall avg fertility | `{s['overall_avg_fertility']}` |\n"
         f"| Overall avg compression | `{s['overall_avg_compression']}` |\n"
         f"| Fidelity failures | `{s['fidelity_failure_count']}` |"
     )
-
-    heatmap = plot_fertility_heatmap(result, model_id.strip())
-    lang_bar = plot_language_fertility_bar(result, model_id.strip()) if "human_languages" in dataset_subset else None
-    scatter = plot_compression_scatter(result, model_id.strip())
-
     rows = []
     for cat, subcats in result.items():
-        if cat.startswith("__") or not isinstance(subcats, dict):
-            continue
+        if cat.startswith("__") or not isinstance(subcats, dict): continue
         for sub, vals in subcats.items():
             if isinstance(vals, dict):
                 rows.append({
@@ -489,39 +598,36 @@ def run_single_eval(model_id: str, tok_type: str, categories: list[str]):
                     "Avg compression": vals.get("avg_compression_ratio", 0),
                     "Fidelity fails": vals.get("fidelity_failures", 0),
                 })
-    df = pd.DataFrame(rows)
-
-    return status, heatmap, lang_bar, scatter, df
+    return (
+        status,
+        plot_fertility_heatmap(result, name),
+        plot_language_fertility_bar(result, name) if "human_languages" in dataset_subset else None,
+        plot_compression_scatter(result, name),
+        pd.DataFrame(rows),
+    )
 
 
 def run_compare_eval(
-    model_a: str, type_a: str,
-    model_b: str, type_b: str,
-    metric: str, categories: list[str],
+    src_a, hub_a, files_a, name_a, tt_a,
+    src_b, hub_b, files_b, name_b, tt_b,
+    metric, categories,
 ):
-    models = [(model_a.strip(), type_a), (model_b.strip(), type_b)]
-    models = [(m, t) for m, t in models if m]
-    if len(models) < 2:
-        return "⚠️ Please enter at least 2 model names.", None, None, None
-
-    tokenizers = {}
-    for model_id, tok_type in models:
+    results_dict = {}
+    for src, hub, files, uname, tt in [
+        (src_a, hub_a, files_a, name_a, tt_a),
+        (src_b, hub_b, files_b, name_b, tt_b),
+    ]:
         try:
-            if tok_type == "HuggingFace (AutoTokenizer)":
-                tokenizers[model_id] = load_hf_tokenizer(model_id)
-            else:
-                tokenizers[model_id] = load_tiktoken(model_id)
+            tok, dname = resolve_tokenizer(src, hub, files, uname, tt)
         except Exception:
-            return f"❌ Failed to load `{model_id}`:\n{traceback.format_exc()}", None, None, None
-
-    dataset_subset = {k: v for k, v in DATASET.items() if k in categories} if categories else DATASET
-
-    results_dict = {}
-    for name, tok in tokenizers.items():
+            return f"❌ Could not load tokenizer:\n```\n{traceback.format_exc()}\n```", None, None, None
+        dataset_subset = {k: v for k, v in DATASET.items() if k in (categories or [])}
+        if not dataset_subset:
+            return "⚠️ Select at least one dataset category.", None, None, None
         try:
-            results_dict[name] = evaluate_tokenizer(tok, dataset_subset)
+            results_dict[dname] = evaluate_tokenizer(tok, dataset_subset)
         except Exception:
-            return f"❌ Evaluation failed for `{name}`:\n{traceback.format_exc()}", None, None, None
+            return f"❌ Eval error for `{dname}`:\n```\n{traceback.format_exc()}\n```", None, None, None
 
     metric_key = {
         "Fertility (lower = better)": "avg_fertility",
@@ -529,9 +635,6 @@ def run_compare_eval(
         "Byte compression": "avg_byte_compression",
     }.get(metric, "avg_fertility")
 
-    cmp_bar = plot_comparison_bar(results_dict, metric_key)
-    fid_bar = plot_fidelity_summary(results_dict)
-
     rows = []
     for name, result in results_dict.items():
         s = result.get("__summary__", {})
@@ -543,140 +646,107 @@ def run_compare_eval(
             "Fidelity failures": s.get("fidelity_failure_count"),
         })
     df = pd.DataFrame(rows).sort_values("Avg fertility")
-
     status = "✅ Comparison complete.\n\n**Leaderboard (lower fertility = better)**\n\n"
     for _, row in df.iterrows():
         status += f"- **{row['Tokenizer']}** — fertility `{row['Avg fertility']}`, failures `{row['Fidelity failures']}`\n"
 
-    return status, cmp_bar, fid_bar, df
-
-
-def tokenize_live(model_id: str, tok_type: str, text: str):
-    if not model_id.strip() or not text.strip():
-        return "Enter a model name and some text above.", ""
-    try:
-        if tok_type == "HuggingFace (AutoTokenizer)":
-            tok = load_hf_tokenizer(model_id.strip())
-        else:
-            tok = load_tiktoken(model_id.strip())
-        ids = tok.encode(text)
-        decoded = tok.decode(ids)
-        fid = "✅ Roundtrip OK" if text.strip() == decoded.strip() else "⚠️ Roundtrip mismatch"
-        info = (
-            f"**Token count:** {len(ids)}  |  "
-            f"**Fertility:** {len(ids)/max(1,len(text.split())):.2f}  |  "
-            f"**Compression:** {len(ids)/max(1,len(text)):.3f}  |  "
-            f"**Fidelity:** {fid}"
-        )
-        ids_str = " ".join(str(i) for i in ids[:100])
-        if len(ids) > 100:
-            ids_str += f" … (+{len(ids)-100} more)"
-        return info, ids_str
-    except Exception:
-        return f"❌ Error:\n{traceback.format_exc()}", ""
+    return (
+        status,
+        plot_comparison_bar(results_dict, metric_key),
+        plot_fidelity_summary(results_dict),
+        df,
+    )
 
 # ─────────────────────────────────────────────────────────────────
 # Gradio UI
 # ────────────────────────────────��────────────────────────────────
 
 CATEGORY_CHOICES = list(DATASET.keys())
-CATEGORY_DEFAULT = CATEGORY_CHOICES
-
-TYPE_CHOICES = ["HuggingFace (AutoTokenizer)", "tiktoken"]
-
-EXAMPLE_HF = ["bert-base-multilingual-cased", "xlm-roberta-base",
-               "google/mt5-base", "facebook/mbart-large-50"]
-EXAMPLE_TIKTOKEN = ["cl100k_base", "o200k_base", "p50k_base"]
 
 with gr.Blocks(title="TokenizerBench", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """# 🤗 TokenizerBench
         Evaluate and compare tokenizers on multilingual text, code, scientific formulas, and edge cases.
-        Built on the [TokenizerBench dataset](https://huggingface.co/datasets).
+        Load from the **Hugging Face Hub**, **upload your own files**, or use a **tiktoken** encoding.
         """
     )
 
     with gr.Tabs():
 
-        # ── Tab 1: Playground ────────────────────────────────────
+        # ── Tab 1: Playground ──────────────────────────────────
         with gr.Tab("🧪 Playground"):
-            gr.Markdown("### Live tokenization — try any text")
+            gr.Markdown("Type or paste any text and see instant tokenization results.")
             with gr.Row():
                 with gr.Column(scale=1):
-                    live_model = gr.Textbox(label="Model name / encoding",
-                                            placeholder="bert-base-multilingual-cased",
-                                            value="bert-base-multilingual-cased")
-                    live_type = gr.Dropdown(TYPE_CHOICES, value=TYPE_CHOICES[0],
-                                            label="Tokenizer type")
+                    pg = tokenizer_source_block()
+                    pg_btn = gr.Button("Tokenize ▶", variant="primary")
                 with gr.Column(scale=2):
-                    live_text = gr.Textbox(
+                    pg_text = gr.Textbox(
                         label="Input text",
                         placeholder="Type or paste anything…",
-                        lines=4,
+                        lines=5,
                         value="The quick brown fox jumps over the lazy dog. 快速的棕色狐狸跳过了懒狗。",
                     )
-            live_btn = gr.Button("Tokenize", variant="primary")
-            live_info = gr.Markdown("Metrics will appear here.")
-            live_ids = gr.Textbox(label="Token IDs", lines=2, interactive=False)
-            live_btn.click(tokenize_live, [live_model, live_type, live_text],
-                           [live_info, live_ids])
+                    pg_info = gr.Markdown("_Results will appear here._")
+                    pg_ids  = gr.Textbox(label="Token IDs", lines=2, interactive=False)
+
+            pg_btn.click(
+                tokenize_live,
+                [pg["source"], pg["hub_id"], pg["uploaded_files"],
+                 pg["upload_name"], pg["tiktoken_enc"], pg_text],
+                [pg_info, pg_ids],
+            )
 
-            gr.Markdown("---\n### Dataset samples — click to load into the text box")
+            gr.Markdown("---\n### Browse dataset samples")
+            gr.Markdown("Click any sample below to load it into the text box above.")
             for cat_key, cat_label in CATEGORY_LABELS.items():
                 with gr.Accordion(cat_label, open=False):
                     for sub, samples in DATASET[cat_key].items():
+                        gr.Markdown(f"**{sub}**")
                         with gr.Row():
                             for s in samples[:3]:
-                                btn = gr.Button(s[:60] + ("…" if len(s) > 60 else ""),
-                                                size="sm")
-                                btn.click(lambda t=s: t, outputs=live_text)
+                                btn = gr.Button(s[:65] + ("…" if len(s) > 65 else ""), size="sm")
+                                btn.click(lambda t=s: t, outputs=pg_text)
 
-        # ── Tab 2: Evaluate ──────────────────────────────────────
+        # ── Tab 2: Evaluate ────────────────────────────────────
         with gr.Tab("📊 Evaluate"):
-            gr.Markdown("### Evaluate a single tokenizer against the full dataset")
+            gr.Markdown("Run a full benchmark on a single tokenizer across all dataset categories.")
             with gr.Row():
                 with gr.Column(scale=1):
-                    eval_model = gr.Textbox(label="Model name / encoding",
-                                            placeholder="xlm-roberta-base",
-                                            value="bert-base-multilingual-cased")
-                    eval_type = gr.Dropdown(TYPE_CHOICES, value=TYPE_CHOICES[0],
-                                            label="Tokenizer type")
-                    eval_cats = gr.CheckboxGroup(
-                        CATEGORY_CHOICES, value=CATEGORY_DEFAULT,
+                    ev = tokenizer_source_block()
+                    ev_cats = gr.CheckboxGroup(
+                        CATEGORY_CHOICES, value=CATEGORY_CHOICES,
                         label="Dataset categories to evaluate",
                     )
-                    eval_btn = gr.Button("Run evaluation", variant="primary")
+                    ev_btn = gr.Button("Run evaluation ▶", variant="primary")
                 with gr.Column(scale=2):
-                    eval_status = gr.Markdown("Results will appear here.")
+                    ev_status = gr.Markdown("_Results will appear here after you click Run evaluation._")
 
-            eval_table = gr.Dataframe(label="Per-subcategory results", wrap=True)
+            ev_table = gr.Dataframe(label="Per-subcategory breakdown", wrap=True)
 
             with gr.Tabs():
                 with gr.Tab("Fertility heatmap"):
-                    eval_heatmap = gr.Image(label="Heatmap", type="pil")
+                    ev_heatmap = gr.Image(type="pil")
                 with gr.Tab("Language fertility bar"):
-                    eval_langbar = gr.Image(label="Language fertility", type="pil")
+                    ev_langbar = gr.Image(type="pil")
                 with gr.Tab("Fertility vs compression"):
-                    eval_scatter = gr.Image(label="Scatter", type="pil")
+                    ev_scatter = gr.Image(type="pil")
 
-            eval_btn.click(
+            ev_btn.click(
                 run_single_eval,
-                [eval_model, eval_type, eval_cats],
-                [eval_status, eval_heatmap, eval_langbar, eval_scatter, eval_table],
+                [ev["source"], ev["hub_id"], ev["uploaded_files"],
+                 ev["upload_name"], ev["tiktoken_enc"], ev_cats],
+                [ev_status, ev_heatmap, ev_langbar, ev_scatter, ev_table],
             )
 
-        # ── Tab 3: Compare ───────────────────────────────────────
+        # ── Tab 3: Compare ─────────────────────────────────────
         with gr.Tab("⚖️ Compare"):
-            gr.Markdown("### Compare two tokenizers side-by-side")
+            gr.Markdown("Compare two tokenizers side-by-side on the same dataset.")
             with gr.Row():
                 with gr.Column():
-                    gr.Markdown("**Tokenizer A**")
-                    cmp_model_a = gr.Textbox(label="Model A", value="bert-base-multilingual-cased")
-                    cmp_type_a = gr.Dropdown(TYPE_CHOICES, value=TYPE_CHOICES[0], label="Type A")
+                    cmp_a = tokenizer_source_block("Tokenizer A — ")
                 with gr.Column():
-                    gr.Markdown("**Tokenizer B**")
-                    cmp_model_b = gr.Textbox(label="Model B", value="xlm-roberta-base")
-                    cmp_type_b = gr.Dropdown(TYPE_CHOICES, value=TYPE_CHOICES[0], label="Type B")
+                    cmp_b = tokenizer_source_block("Tokenizer B — ")
 
             with gr.Row():
                 cmp_metric = gr.Dropdown(
@@ -685,36 +755,46 @@ with gr.Blocks(title="TokenizerBench", theme=gr.themes.Soft()) as demo:
                     label="Comparison metric",
                 )
                 cmp_cats = gr.CheckboxGroup(
-                    CATEGORY_CHOICES, value=CATEGORY_DEFAULT,
+                    CATEGORY_CHOICES, value=CATEGORY_CHOICES,
                     label="Dataset categories",
                 )
 
-            cmp_btn = gr.Button("Compare", variant="primary")
-            cmp_status = gr.Markdown("Results will appear here.")
-            cmp_table = gr.Dataframe(label="Summary leaderboard", wrap=True)
+            cmp_btn = gr.Button("Compare ▶", variant="primary")
+            cmp_status = gr.Markdown("_Results will appear here._")
+            cmp_table  = gr.Dataframe(label="Leaderboard", wrap=True)
 
             with gr.Tabs():
                 with gr.Tab("Category comparison bar"):
-                    cmp_bar_img = gr.Image(label="Grouped bar", type="pil")
+                    cmp_bar_img = gr.Image(type="pil")
                 with gr.Tab("Fidelity failures"):
-                    cmp_fid_img = gr.Image(label="Fidelity", type="pil")
+                    cmp_fid_img = gr.Image(type="pil")
 
             cmp_btn.click(
                 run_compare_eval,
-                [cmp_model_a, cmp_type_a, cmp_model_b, cmp_type_b, cmp_metric, cmp_cats],
+                [
+                    cmp_a["source"], cmp_a["hub_id"], cmp_a["uploaded_files"],
+                    cmp_a["upload_name"], cmp_a["tiktoken_enc"],
+                    cmp_b["source"], cmp_b["hub_id"], cmp_b["uploaded_files"],
+                    cmp_b["upload_name"], cmp_b["tiktoken_enc"],
+                    cmp_metric, cmp_cats,
+                ],
                 [cmp_status, cmp_bar_img, cmp_fid_img, cmp_table],
             )
 
     gr.Markdown(
         """---
-        **Dataset categories:** Human languages (8 languages) · Programming languages (Python, JS, SQL, Rust) · Scientific formulas (algebra, calculus, physics, stats) · Edge cases (whitespace, long tokens, mixed scripts)
+        **Metrics explained** — Fertility = tokens/word (lower = better, ≥4 = poor) · Compression = tokens/char · Fidelity = encode→decode must reproduce original text exactly
+
+        **Upload guide**
 
-        **Metrics explained:**
-        - **Fertility** — tokens per word (lower = more efficient; ≥4 = poor coverage)
-        - **Compression ratio** — tokens per character
-        - **Fidelity** — roundtrip encode→decode produces identical text (must be 1.0)
+        | File(s) to upload | Tokenizer type |
+        |-------------------|----------------|
+        | `tokenizer.json` | Any HuggingFace fast tokenizer (BERT, RoBERTa, GPT-2, LLaMA…) |
+        | `tokenizer.json` + `tokenizer_config.json` + `vocab.txt` | Full HF tokenizer folder |
+        | `vocab.json` + `merges.txt` | BPE tokenizer (GPT-2 style) |
+        | `*.model` | SentencePiece (T5, mT5, XLM-R, mBERT…) |
         """
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

requirements.txt

@@ -1,4 +1,4 @@
-gradio>=4.0.0
+gradio>=6.11.0
 transformers>=4.38.0
 tiktoken>=0.6.0
 torch>=2.0.0