Create app.py

app.py

@@ -0,0 +1,412 @@
+from __future__ import annotations
+
+import functools
+import re
+import time
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import gradio as gr
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+MODEL_ID = "fromziro/JetonCount"
+TOKENIZER_ID = "AxiomicLabs/GPT-X2-125M"
+
+# Fill these from your training artifact if you want exact parity.
+FEATURE_MEAN = None
+FEATURE_STD = None
+TARGET_OFFSET = 0.0
+
+DEFAULT_VOCAB_SIZE = 32_000
+
+TEXT = """
+orfffffffdhdeheorihrhehryh3ureh {
+  "actual_token_count": 6,
+  "prediction": -5.110997403079409e-09,
+  "model_latency_ms": 0.35677890000033585,
+  "tokenizer_latency_ms": 0.10765600000013364,
+  "model_id": "fromziro/JetonCount",
+  "tokenizer_id": "fromziro/Er-Tiny-1.3M",
+  "vocab_size": 2564,
+  "features": {
+    "chars": 19.0,
+    "words": 4.0,
+    "avg_chars_per_word": 3.75,
+    "punctuation_ratio": 0.05263157894736842,
+    "symbol_ratio": 0.0,
+    "longest_word_chars": 4.0,
+    "vocab_size": 2564.0
+  }
+}{
+  "actual_token_count": 6,
+  "prediction": -5.110997403079409e-09,
+  "model_latency_ms": 0.35677890000033585,
+  "tokenizer_latency_ms": 0.10765600000013364,
+  "model_id": "fromziro/JetonCount",
+  "tokenizer_id": "fromziro/Er-Tiny-1.3M",
+  "vocab_size": 2564,
+  "features": {
+    "chars": 19.0,
+    "words": 4.0,
+    "avg_chars_per_word": 3.75,
+    "punctuation_ratio": 0.05263157894736842,
+    "symbol_ratio": 0.0,
+    "longest_word_chars": 4.0,
+    "vocab_size": 2564.0
+  }
+}
+"""
+
+TOKENIZER_ROUNDS = 100
+MODEL_ROUNDS = 1000
+
+PUNCTUATION_CHARS = set(r""".,!?;:'"`~@#$%^&*()-_=+[]{}<>/\|""")
+SYMBOL_CHARS = set(r"""@#$%^&*()-_=+[]{}<>/\|~`""")
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+@dataclass
+class TextStats:
+    chars: float
+    words: float
+    avg_chars_per_word: float
+    punctuation_ratio: float
+    symbol_ratio: float
+    longest_word_chars: float
+    vocab_size: float
+
+
+def _safe_int(value: object, default: int) -> int:
+    try:
+        return int(float(value))
+    except Exception:
+        return default
+
+
+def _clean_tokenizer_id(tokenizer_id: str | None) -> str:
+    return (tokenizer_id or "").strip()
+
+
+@functools.lru_cache(maxsize=2)
+def load_model(model_id: str) -> AutoModel:
+    model = AutoModel.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+    )
+    model.eval()
+    model.to(DEVICE)
+    return model
+
+
+@functools.lru_cache(maxsize=32)
+def load_tokenizer(tokenizer_id: str):
+    return AutoTokenizer.from_pretrained(tokenizer_id, use_fast=True)
+
+
+def get_tokenizer_vocab_size(tokenizer) -> int:
+    vocab_size = getattr(tokenizer, "vocab_size", None)
+    if isinstance(vocab_size, int) and vocab_size > 0:
+        return vocab_size
+    try:
+        return int(len(tokenizer))
+    except Exception:
+        return DEFAULT_VOCAB_SIZE
+
+
+def compute_text_stats(text: str, vocab_size: int) -> TextStats:
+    chars = len(text)
+    words_list = re.findall(r"\b\w+\b", text, flags=re.UNICODE)
+    words = len(words_list)
+
+    total_word_chars = sum(len(w) for w in words_list)
+    avg_chars_per_word = (total_word_chars / words) if words else 0.0
+    longest_word_chars = max((len(w) for w in words_list), default=0)
+
+    if chars:
+        punctuation_count = sum(1 for ch in text if ch in PUNCTUATION_CHARS)
+        symbol_count = sum(1 for ch in text if ch in SYMBOL_CHARS)
+        punctuation_ratio = punctuation_count / chars
+        symbol_ratio = symbol_count / chars
+    else:
+        punctuation_ratio = 0.0
+        symbol_ratio = 0.0
+
+    return TextStats(
+        chars=float(chars),
+        words=float(words),
+        avg_chars_per_word=float(avg_chars_per_word),
+        punctuation_ratio=float(punctuation_ratio),
+        symbol_ratio=float(symbol_ratio),
+        longest_word_chars=float(longest_word_chars),
+        vocab_size=float(vocab_size),
+    )
+
+
+def build_feature_tensor(stats: TextStats) -> torch.Tensor:
+    base = torch.tensor(
+        [
+            stats.chars,
+            stats.words,
+            stats.avg_chars_per_word,
+            stats.punctuation_ratio,
+            stats.symbol_ratio,
+            stats.longest_word_chars,
+            stats.vocab_size,
+        ],
+        dtype=torch.float32,
+    )
+
+    chars, words, avg_chars_per_word, punctuation_ratio, symbol_ratio, longest_word_chars, vocab_size = base
+    eps = 1e-6
+
+    extra = torch.tensor(
+        [
+            chars / max(words.item(), 1.0),
+            words / max(chars.item(), 1.0),
+            torch.log1p(torch.clamp(chars, min=0.0)).item(),
+            torch.log1p(torch.clamp(words, min=0.0)).item(),
+            torch.log1p(torch.clamp(vocab_size, min=0.0)).item(),
+            (chars * punctuation_ratio).item(),
+            (chars * symbol_ratio).item(),
+            (words * avg_chars_per_word).item(),
+            (words * punctuation_ratio).item(),
+            (longest_word_chars * punctuation_ratio).item(),
+            ((avg_chars_per_word + longest_word_chars) * (1.0 + punctuation_ratio + symbol_ratio)).item(),
+            ((chars + eps) * (punctuation_ratio + symbol_ratio + eps)).item(),
+        ],
+        dtype=torch.float32,
+    )
+
+    return torch.cat([base, extra], dim=0)
+
+
+def standardize_features(x: torch.Tensor) -> torch.Tensor:
+    if FEATURE_MEAN is None or FEATURE_STD is None:
+        return x
+
+    mean = torch.tensor(FEATURE_MEAN, dtype=x.dtype, device=x.device)
+    std = torch.tensor(FEATURE_STD, dtype=x.dtype, device=x.device)
+
+    safe_std = torch.where(torch.isfinite(std) & (std != 0), std, torch.ones_like(std))
+    safe_mean = torch.where(torch.isfinite(mean), mean, torch.zeros_like(mean))
+    return (x - safe_mean) / safe_std
+
+
+def benchmark_tokenizer(tokenizer, text: str, rounds: int = 100) -> Tuple[int, float]:
+    _ = tokenizer(text, add_special_tokens=False)
+    start = time.perf_counter()
+    actual_count = 0
+    for _ in range(rounds):
+        ids = tokenizer(text, add_special_tokens=False).input_ids
+        actual_count = len(ids)
+    elapsed_ms = (time.perf_counter() - start) * 1000.0 / rounds
+    return actual_count, elapsed_ms
+
+
+@torch.inference_mode()
+def benchmark_model(model, feature_tensor: torch.Tensor, rounds: int = 1000) -> Tuple[float, float]:
+    x = standardize_features(feature_tensor).unsqueeze(0).to(DEVICE)
+
+    _ = model(input_features=x)
+
+    start = time.perf_counter()
+    pred = 0.0
+    for _ in range(rounds):
+        out = model(input_features=x)
+        pred = float(out.logits.squeeze().item())
+    elapsed_ms = (time.perf_counter() - start) * 1000.0 / rounds
+    return pred, elapsed_ms
+
+
+def update_vocab_lock(tokenizer_repo_id: str):
+    repo_id = _clean_tokenizer_id(tokenizer_repo_id)
+
+    if not repo_id:
+        return (
+            gr.update(value=DEFAULT_VOCAB_SIZE, interactive=True),
+            f"Using manual vocab size. Default tokenizer: `{TOKENIZER_ID}`.",
+            None,
+        )
+
+    try:
+        tokenizer = load_tokenizer(repo_id)
+        vocab_size = get_tokenizer_vocab_size(tokenizer)
+        return (
+            gr.update(value=vocab_size, interactive=False),
+            f"Tokenizer locked to `{repo_id}` with vocab size `{vocab_size}`.",
+            vocab_size,
+        )
+    except Exception as exc:
+        return (
+            gr.update(interactive=True),
+            f"Could not load tokenizer `{repo_id}`: {exc}",
+            None,
+        )
+
+
+def run_benchmark(
+    text: str,
+    vocab_size_value: object,
+    tokenizer_repo_id: str,
+    locked_vocab_size: Optional[int],
+):
+    text = text or ""
+    repo_id = _clean_tokenizer_id(tokenizer_repo_id)
+
+    try:
+        if repo_id:
+            tokenizer = load_tokenizer(repo_id)
+            resolved_vocab_size = (
+                locked_vocab_size if locked_vocab_size is not None else get_tokenizer_vocab_size(tokenizer)
+            )
+            tokenizer_source = repo_id
+            tokenizer_locked = True
+        else:
+            tokenizer = load_tokenizer(TOKENIZER_ID)
+            resolved_vocab_size = _safe_int(vocab_size_value, DEFAULT_VOCAB_SIZE)
+            tokenizer_source = TOKENIZER_ID
+            tokenizer_locked = False
+
+        resolved_vocab_size = int(resolved_vocab_size)
+        stats = compute_text_stats(text, resolved_vocab_size)
+        feature_tensor = build_feature_tensor(stats)
+
+        actual_count, tokenizer_latency_ms = benchmark_tokenizer(tokenizer, text, rounds=TOKENIZER_ROUNDS)
+        model = load_model(MODEL_ID)
+        prediction, model_latency_ms = benchmark_model(model, feature_tensor, rounds=MODEL_ROUNDS)
+
+        result = {
+            "actual_token_count": actual_count,
+            "prediction": prediction + TARGET_OFFSET,
+            "model_latency_ms": model_latency_ms,
+            "tokenizer_latency_ms": tokenizer_latency_ms,
+            "model_id": MODEL_ID,
+            "tokenizer_id": tokenizer_source,
+            "vocab_size": resolved_vocab_size,
+            "tokenizer_locked": tokenizer_locked,
+            "features": {
+                "chars": stats.chars,
+                "words": stats.words,
+                "avg_chars_per_word": stats.avg_chars_per_word,
+                "punctuation_ratio": stats.punctuation_ratio,
+                "symbol_ratio": stats.symbol_ratio,
+                "longest_word_chars": stats.longest_word_chars,
+                "vocab_size": stats.vocab_size,
+            },
+        }
+
+        status = (
+            f"Benchmarked `{tokenizer_source}` on `{DEVICE.type}`. "
+            f"Vocab size used: `{resolved_vocab_size}`."
+        )
+        return result, status
+
+    except Exception as exc:
+        return None, f"Benchmark failed: {exc}"
+
+
+CSS = """
+:root {
+    --background-fill-primary: #000000 !important;
+    --background-fill-secondary: #0b0b0b !important;
+    --block-background-fill: #0b0b0b !important;
+    --body-text-color: #f4f4f5 !important;
+    --link-text-color: #ffffff !important;
+    --button-primary-background-fill: #ffffff !important;
+    --button-primary-text-color: #000000 !important;
+}
+
+body, .gradio-container {
+    background: #000 !important;
+    color: #f4f4f5 !important;
+}
+
+.gradio-container {
+    max-width: 1200px !important;
+}
+
+h1, h2, h3, h4, p, label, span, div {
+    color: #f4f4f5 !important;
+}
+
+textarea, input, .wrap, .prose, .block, .block-container, .panel, .container {
+    background-color: #090909 !important;
+    color: #f4f4f5 !important;
+}
+
+textarea, input {
+    border: 1px solid #2b2b2b !important;
+}
+
+button {
+    border-radius: 14px !important;
+}
+
+#status_box {
+    border: 1px solid #222 !important;
+    background: #080808 !important;
+    padding: 12px !important;
+    border-radius: 16px !important;
+}
+"""
+
+
+with gr.Blocks(
+    theme=gr.themes.Monochrome(),
+    css=CSS,
+    title="JetonCount Gradio Bench",
+) as demo:
+    gr.Markdown("# JetonCount Gradio Bench")
+    gr.Markdown(
+        "Paste text, choose a vocab size, and optionally provide a tokenizer HF repo ID. "
+        "When a tokenizer repo ID is present, the vocab field locks to that tokenizer's vocab size for a fair benchmark."
+    )
+
+    locked_vocab_state = gr.State(None)
+
+    with gr.Row():
+        text_in = gr.Textbox(
+            label="Text",
+            value=TEXT,
+            lines=16,
+            placeholder="Paste text here...",
+            interactive=True,
+        )
+
+    with gr.Row():
+        vocab_in = gr.Number(
+            label="Vocab size",
+            value=DEFAULT_VOCAB_SIZE,
+            interactive=True,
+        )
+        tokenizer_in = gr.Textbox(
+            label="Tokenizer HF Repo ID (optional)",
+            value="",
+            placeholder="e.g. AxiomicLabs/GPT-X2-125M",
+            interactive=True,
+        )
+
+    status = gr.Markdown(
+        value=f"Using manual vocab size. Default tokenizer: `{TOKENIZER_ID}`.",
+        elem_id="status_box",
+    )
+
+    run_btn = gr.Button("Benchmark", variant="primary")
+    result_out = gr.JSON(label="Result")
+
+    tokenizer_in.blur(
+        fn=update_vocab_lock,
+        inputs=tokenizer_in,
+        outputs=[vocab_in, status, locked_vocab_state],
+    )
+
+    run_btn.click(
+        fn=run_benchmark,
+        inputs=[text_in, vocab_in, tokenizer_in, locked_vocab_state],
+        outputs=[result_out, status],
+    )
+
+if __name__ == "__main__":
+    demo.launch()