token-atlas-engine 0.3.0__py3-none-any.whl

token_atlas/__init__.py

@@ -0,0 +1,38 @@
+"""TokenAtlas: Universal Multilingual Tokenizer Benchmark & Global Inflation Atlas."""
+
+__version__ = "0.3.0"
+__author__ = "Prajwal Pujari"
+
+from token_atlas.models import (
+    TokenizationResult,
+    LanguageComparison,
+    BenchmarkReport,
+    TokenizerSpec,
+)
+from token_atlas.metrics import (
+    calculate_inflation_index,
+    calculate_bytes_per_token,
+    calculate_effective_context,
+    calculate_cost_multiplier,
+)
+from token_atlas.router import recommend_model, ModelRecommendation
+from token_atlas.minify import minify_prompt, MinificationResult
+from token_atlas.domains import evaluate_domain, DomainBenchmarkReport
+
+__all__ = [
+    "__version__",
+    "TokenizationResult",
+    "LanguageComparison",
+    "BenchmarkReport",
+    "TokenizerSpec",
+    "calculate_inflation_index",
+    "calculate_bytes_per_token",
+    "calculate_effective_context",
+    "calculate_cost_multiplier",
+    "recommend_model",
+    "ModelRecommendation",
+    "minify_prompt",
+    "MinificationResult",
+    "evaluate_domain",
+    "DomainBenchmarkReport",
+]

token_atlas/bpe.py

@@ -0,0 +1,98 @@
+"""Educational First-Principles Byte-Pair Encoding (BPE) Implementation.
+
+Demonstrates how BPE merge tables are constructed from raw UTF-8 bytes and
+why under-represented scripts suffer from severe byte-fallback fragmentation.
+"""
+
+from typing import NamedTuple
+from collections import Counter
+
+
+class BPETrainResult(NamedTuple):
+    vocab: dict[int, bytes]
+    merges: list[tuple[bytes, bytes]]
+    token_ids: list[int]
+
+
+def get_stats(ids: list[int]) -> Counter[tuple[int, int]]:
+    """Count frequency of consecutive token id pairs."""
+    counts: Counter[tuple[int, int]] = Counter()
+    for p in zip(ids[:-1], ids[1:]):
+        counts[p] += 1
+    return counts
+
+
+def merge_ids(ids: list[int], pair: tuple[int, int], new_id: int) -> list[int]:
+    """Replace all occurrences of `pair` in `ids` with `new_id`."""
+    new_ids: list[int] = []
+    i = 0
+    n = len(ids)
+    while i < n:
+        if i < n - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
+            new_ids.append(new_id)
+            i += 2
+        else:
+            new_ids.append(ids[i])
+            i += 1
+    return new_ids
+
+
+def train_toy_bpe(text: str, num_merges: int = 20) -> BPETrainResult:
+    """Train a first-principles BPE tokenizer from raw UTF-8 text bytes.
+    
+    Initializes vocabulary with 256 byte tokens (0..255).
+    Iteratively merges the most frequent adjacent pair.
+    """
+    raw_bytes = list(text.encode("utf-8"))
+    ids = list(raw_bytes)
+
+    # Base vocabulary: 256 individual UTF-8 bytes
+    vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
+    merges: list[tuple[bytes, bytes]] = []
+
+    for i in range(num_merges):
+        stats = get_stats(ids)
+        if not stats:
+            break
+        # Pick the most frequent pair
+        pair, count = stats.most_common(1)[0]
+        if count < 2:
+            # No repeated pairs left to merge
+            break
+        
+        new_id = 256 + i
+        ids = merge_ids(ids, pair, new_id)
+        
+        piece_a = vocab[pair[0]]
+        piece_b = vocab[pair[1]]
+        merged_piece = piece_a + piece_b
+        vocab[new_id] = merged_piece
+        merges.append((piece_a, piece_b))
+
+    return BPETrainResult(vocab=vocab, merges=merges, token_ids=ids)
+
+
+def explain_fragmentation(text: str) -> dict[str, any]:
+    """Diagnose UTF-8 byte fragmentation for an input string."""
+    encoded_bytes = text.encode("utf-8")
+    byte_len = len(encoded_bytes)
+    char_len = len(text)
+    
+    char_breakdown = []
+    for ch in text[:10]:
+        ch_bytes = ch.encode("utf-8")
+        char_breakdown.append({
+            "char": ch,
+            "unicode": f"U+{ord(ch):04X}",
+            "byte_count": len(ch_bytes),
+            "raw_hex": " ".join([f"0x{b:02X}" for b in ch_bytes]),
+        })
+
+    return {
+        "text_sample": text[:30],
+        "character_count": char_len,
+        "utf8_byte_count": byte_len,
+        "bytes_per_char": byte_len / max(1, char_len),
+        "sample_breakdown": char_breakdown,
+    }
+

token_atlas/cli.py

@@ -0,0 +1,402 @@
+"""TokenAtlas CLI: Universal Multilingual Tokenizer Benchmark."""
+
+import json
+from pathlib import Path
+from typing import Optional
+import typer
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+
+from token_atlas import __version__
+from token_atlas.corpora import SUPPORTED_LANGUAGES
+from token_atlas.registry import TokenizerRegistry
+from token_atlas.engine import compare_single_text, run_benchmark
+from token_atlas.web import generate_web_dashboard
+
+app = typer.Typer(
+    name="token-atlas",
+    help="TokenAtlas: Universal Multilingual Tokenizer Benchmark & Global Inflation Atlas",
+    add_completion=False,
+)
+console = Console()
+
+
+def version_callback(value: bool):
+    if value:
+        console.print(f"[bold cyan]TokenAtlas[/bold cyan] version [bold green]{__version__}[/bold green]")
+        raise typer.Exit()
+
+
+@app.callback()
+def callback(
+    version: Optional[bool] = typer.Option(
+        None, "--version", "-v", help="Show version and exit", callback=version_callback, is_eager=True
+    ),
+):
+    pass
+
+
+def main():
+    app()
+
+
+@app.command("list-tokenizers")
+def list_tokenizers(
+    category: Optional[str] = typer.Option(
+        None, "--category", "-c", help="Filter by category (general, multilingual, code, reasoning, slm)"
+    )
+):
+    """List all supported frontier tokenizer architectures and metadata."""
+    reg = TokenizerRegistry()
+    title_suffix = f" (Category: {category})" if category else ""
+    table = Table(title=f"[bold cyan]Supported Tokenizers in TokenAtlas{title_suffix}[/bold cyan]")
+    table.add_column("ID", style="bold green", no_wrap=True)
+    table.add_column("Display Name", style="white")
+    table.add_column("Category", style="blue", no_wrap=True)
+    table.add_column("Vocab Size", justify="right", style="cyan", no_wrap=True)
+    table.add_column("Context Window", justify="right", style="magenta", no_wrap=True)
+    table.add_column("Backend", style="yellow", no_wrap=True)
+    table.add_column("Key Strengths & Compression Profile", style="dim white")
+
+    for spec in reg.list_tokenizers(category=category):
+        table.add_row(
+            spec.id,
+            spec.name,
+            spec.category.upper(),
+            f"{spec.vocab_size:,}",
+            f"{spec.context_window:,} tok",
+            spec.backend,
+            spec.notes or "—",
+        )
+    console.print(table)
+
+
+@app.command("list-languages")
+def list_languages():
+    """List all supported language evaluation benchmarks and language families."""
+    table = Table(title="[bold cyan]Supported Languages (FLORES-200 Parallel Corpus)[/bold cyan]")
+    table.add_column("Code", style="bold green", justify="center")
+    table.add_column("Language", style="white")
+    table.add_column("Family", style="magenta")
+    table.add_column("Script", style="yellow")
+
+    for meta in SUPPORTED_LANGUAGES.values():
+        table.add_row(meta.code, meta.name, meta.family, meta.script)
+    console.print(table)
+
+
+@app.command("compare")
+def compare_cmd(
+    text: Optional[str] = typer.Option(None, "--text", "-t", help="Text to tokenize across models"),
+    file: Optional[Path] = typer.Option(None, "--file", "-f", help="Path to text file to tokenize"),
+    tokenizers: str = typer.Option("gpt-4o,gpt-4", "--tokenizers", "-m", help="Comma-separated tokenizer IDs"),
+    language: str = typer.Option("en", "--language", "-l", help="Language code (e.g. en, ja, hi, de)"),
+):
+    """Compare how multiple tokenizers encode a single text string or file."""
+    if file is not None:
+        if not file.exists():
+            console.print(f"[bold red]Error: File not found at {file}[/bold red]")
+            raise typer.Exit(1)
+        raw_text = file.read_text(encoding="utf-8")
+    elif text is not None:
+        raw_text = text
+    else:
+        console.print("[bold red]Error: Either --text or --file must be specified.[/bold red]")
+        raise typer.Exit(1)
+
+    tids = [t.strip() for t in tokenizers.split(",") if t.strip()]
+    comp = compare_single_text(raw_text, tids, language=language)
+
+    preview = raw_text[:300] + "..." if len(raw_text) > 300 else raw_text
+    console.print(Panel(f"[bold white]{preview}[/bold white]", title=f"[bold cyan]Input Text ({comp.language_name})[/bold cyan]"))
+
+    table = Table(title="[bold]Tokenizer Efficiency Comparison[/bold]")
+    table.add_column("Tokenizer", style="bold green")
+    table.add_column("Tokens", justify="right", style="cyan")
+    table.add_column("Bytes", justify="right", style="white")
+    table.add_column("Bytes / Token", justify="right", style="magenta")
+    table.add_column("First Token Pieces (Sample)", style="dim")
+
+    for res in comp.results:
+        pieces_display = " ".join([f"[{p}]" for p in res.token_pieces[:10]])
+        table.add_row(
+            res.tokenizer_id,
+            str(res.token_count),
+            str(res.byte_count),
+            f"{res.bytes_per_token:.2f}",
+            pieces_display,
+        )
+    console.print(table)
+
+
+@app.command("bench")
+def bench_cmd(
+    tokenizers: str = typer.Option("gpt-4o,gpt-4", "--tokenizers", "-m", help="Comma-separated tokenizer IDs"),
+    languages: str = typer.Option("all", "--languages", "-l", help="Comma-separated language codes or 'all'"),
+    samples: int = typer.Option(5, "--samples", "-n", help="Number of parallel sentences to evaluate"),
+    custom_corpus: Optional[Path] = typer.Option(None, "--custom-corpus", help="Path to JSON file containing list of parallel sentence dicts"),
+    json_out: Optional[Path] = typer.Option(None, "--json-out", help="Path to export raw JSON report"),
+    md_out: Optional[Path] = typer.Option(None, "--md-out", help="Path to export markdown summary"),
+):
+    """Run cross-lingual benchmark across parallel sentences and compute inflation indices."""
+    tids = [t.strip() for t in tokenizers.split(",") if t.strip()]
+    
+    if languages.strip().lower() == "all":
+        langs = list(SUPPORTED_LANGUAGES.keys())
+    else:
+        langs = [lang.strip() for lang in languages.split(",") if lang.strip()]
+
+    corpus_data = None
+    source_name = None
+    if custom_corpus is not None:
+        if not custom_corpus.exists():
+            console.print(f"[bold red]Error: Custom corpus file not found at {custom_corpus}[/bold red]")
+            raise typer.Exit(1)
+        with open(custom_corpus, "r", encoding="utf-8") as f:
+            corpus_data = json.load(f)
+        source_name = f"Custom Corpus ({custom_corpus.name})"
+
+    with console.status("[bold cyan]Running parallel multilingual benchmark...[/bold cyan]"):
+        report = run_benchmark(
+            tokenizer_ids=tids,
+            language_codes=langs,
+            sample_limit=samples,
+            custom_corpus=corpus_data,
+            corpora_source=source_name,
+        )
+
+    for bench in report.results:
+        table = Table(title=f"[bold cyan]{bench.tokenizer_name} (Vocab: {bench.vocab_size:,})[/bold cyan]")
+        table.add_column("Lang", style="bold green", justify="center")
+        table.add_column("Language Name", style="white")
+        table.add_column("Family", style="dim")
+        table.add_column("Tokens", justify="right", style="white")
+        table.add_column("Inflation (τ)", justify="right", style="bold yellow")
+        table.add_column("Effective Context", justify="right", style="cyan")
+        table.add_column("Cost Multiplier", justify="right", style="magenta")
+        table.add_column("Bytes/Tok", justify="right", style="dim")
+
+        for code, m in bench.metrics.items():
+            tau_style = "bold green" if m.inflation_index <= 1.2 else ("bold yellow" if m.inflation_index <= 2.0 else "bold red")
+            table.add_row(
+                m.language_code,
+                m.language_name,
+                m.family,
+                f"{m.token_count:,}",
+                f"[{tau_style}]{m.inflation_index:.2f}x[/{tau_style}]",
+                f"{m.effective_context:,}",
+                f"{m.cost_multiplier:.2f}x",
+                f"{m.bytes_per_token:.2f}",
+            )
+        console.print(table)
+        console.print()
+
+    if json_out:
+        json_out.parent.mkdir(parents=True, exist_ok=True)
+        with open(json_out, "w", encoding="utf-8") as f:
+            json.dump(report.model_dump(), f, indent=2)
+        console.print(f"[bold green]✓ Exported JSON report to {json_out}[/bold green]")
+
+    if md_out:
+        md_out.parent.mkdir(parents=True, exist_ok=True)
+        with open(md_out, "w", encoding="utf-8") as f:
+            f.write(f"# TokenAtlas Multilingual Benchmark Report\n\n")
+            f.write(f"- **Evaluated:** {report.timestamp}\n")
+            f.write(f"- **Source:** {report.corpora_source}\n")
+            f.write(f"- **Sample Sentences:** {samples}\n\n")
+            for bench in report.results:
+                f.write(f"## {bench.tokenizer_name} (Vocab: {bench.vocab_size:,})\n\n")
+                f.write("| Language | Code | Family | Tokens | Inflation Index (τ) | Effective Context | Cost Multiplier | Bytes / Token |\n")
+                f.write("| :--- | :---: | :--- | :---: | :---: | :---: | :---: | :---: |\n")
+                for m in bench.metrics.values():
+                    f.write(f"| {m.language_name} | `{m.language_code}` | {m.family} | {m.token_count:,} | **{m.inflation_index:.2f}x** | {m.effective_context:,} | {m.cost_multiplier:.2f}x | {m.bytes_per_token:.2f} |\n")
+                f.write("\n")
+        console.print(f"[bold green]✓ Exported Markdown summary to {md_out}[/bold green]")
+
+
+@app.command("export-web")
+def export_web_cmd(
+    benchmark_json: Path = typer.Option(
+        Path("benchmarks/flores200_benchmark.json"),
+        "--benchmark-json",
+        "-b",
+        help="Path to precomputed benchmark JSON report",
+    ),
+    out_dir: Path = typer.Option(
+        Path("docs"),
+        "--out-dir",
+        "-o",
+        help="Output directory for static web dashboard",
+    ),
+):
+    """Generate standalone interactive Web Atlas dashboard for GitHub Pages / HF Space."""
+    if not benchmark_json.exists():
+        console.print(f"[bold red]Error: Benchmark report not found at {benchmark_json}[/bold red]")
+        console.print("Run 'token-atlas bench --languages all --json-out benchmarks/flores200_benchmark.json' first.")
+        raise typer.Exit(1)
+
+    with open(benchmark_json, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    with console.status(f"[bold cyan]Compiling interactive Web Atlas into {out_dir}...[/bold cyan]"):
+        out_file = generate_web_dashboard(data, out_dir=out_dir)
+
+    console.print(f"[bold green]✓ Successfully generated Web Atlas dashboard at {out_file}[/bold green]")
+    console.print(f"[dim]Deployable directly to GitHub Pages (docs/) or Hugging Face Spaces (static HTML).[/dim]")
+
+
+@app.command("explain-bytes")
+def explain_bytes_cmd(
+    text: str = typer.Option(..., "--text", "-t", help="Text to inspect UTF-8 byte breakdown"),
+):
+    """Diagnose UTF-8 byte layout and explain tokenizer byte-fallback fragmentation."""
+    from token_atlas.bpe import explain_fragmentation
+    diag = explain_fragmentation(text)
+
+    console.print(Panel(f"[bold white]{text}[/bold white]", title="[bold cyan]UTF-8 Byte Fragmentation Diagnostic[/bold cyan]"))
+    console.print(f"Characters: [bold cyan]{diag['character_count']}[/bold cyan] | UTF-8 Bytes: [bold magenta]{diag['utf8_byte_count']}[/bold magenta] | Density: [bold yellow]{diag['bytes_per_char']:.2f} bytes/char[/bold yellow]\n")
+
+    table = Table(title="[bold]Per-Character UTF-8 Byte Breakdown[/bold]")
+    table.add_column("Char", style="bold green", justify="center")
+    table.add_column("Unicode", style="cyan")
+    table.add_column("Bytes", justify="right", style="magenta")
+    table.add_column("Raw Hex Bytes", style="dim")
+
+    for item in diag["sample_breakdown"]:
+        table.add_row(item["char"], item["unicode"], str(item["byte_count"]), item["raw_hex"])
+    console.print(table)
+
+
+@app.command("bench-domain")
+def bench_domain(
+    domain: str = typer.Option("all", "--domain", "-d", help="Domain to benchmark: code, indentation, tools, math, or all"),
+    tokenizers: str = typer.Option("gpt-4o,gpt-4", "--tokenizers", "-t", help="Comma-separated tokenizer IDs"),
+    json_out: Optional[Path] = typer.Option(None, "--json-out", help="Path to export raw benchmark JSON"),
+):
+    """Benchmark multi-domain tokenization efficiency (Code, Indentation, Tool Calling, Math)."""
+    from token_atlas.domains import evaluate_domain
+    tok_list = [t.strip() for t in tokenizers.split(",") if t.strip()]
+
+    with console.status(f"[bold cyan]Evaluating domain '{domain}' across {len(tok_list)} tokenizers...[/bold cyan]"):
+        report = evaluate_domain(domain=domain, tokenizer_ids=tok_list)
+
+    table = Table(title=f"[bold cyan]Multi-Domain Tokenization Benchmark: {domain.upper()}[/bold cyan]")
+    table.add_column("Domain", style="magenta", no_wrap=True)
+    table.add_column("Sample / Task", style="white")
+    table.add_column("Tokenizer", style="bold green", no_wrap=True)
+    table.add_column("Tokens", justify="right", style="cyan", no_wrap=True)
+    table.add_column("Bytes", justify="right", style="dim", no_wrap=True)
+    table.add_column("BPT", justify="right", style="yellow", no_wrap=True)
+    table.add_column("Rel. Tax", justify="right", style="bold red", no_wrap=True)
+
+    for row in report.rows:
+        rel_style = "bold red" if row.relative_inflation > 1.25 else ("bold green" if row.relative_inflation < 0.95 else "white")
+        table.add_row(
+            row.domain.upper(),
+            row.sample_title,
+            row.tokenizer_id,
+            f"{row.token_count:,}",
+            f"{row.byte_count:,}",
+            f"{row.bytes_per_token:.2f}",
+            f"[{rel_style}]{row.relative_inflation:.2f}x[/{rel_style}]",
+        )
+    console.print(table)
+
+    if report.summary_metrics:
+        console.print()
+        if "indentation_tax_ratio" in report.summary_metrics:
+            summary_table = Table(title="[bold yellow]Indentation Tax Summary (4-Space vs Minified)[/bold yellow]")
+            summary_table.add_column("Tokenizer", style="bold green")
+            summary_table.add_column("Indentation Inflation Multiplier", style="bold red")
+            for tok, ratio in report.summary_metrics["indentation_tax_ratio"].items():
+                summary_table.add_row(tok, f"{ratio:.2f}x")
+            console.print(summary_table)
+
+        if "tool_schema_whitespace_tax" in report.summary_metrics:
+            summary_table = Table(title="[bold yellow]Tool-Calling Schema Spacing Tax (Verbose vs Compact)[/bold yellow]")
+            summary_table.add_column("Tokenizer", style="bold green")
+            summary_table.add_column("Schema Spacing Multiplier", style="bold red")
+            for tok, ratio in report.summary_metrics["tool_schema_whitespace_tax"].items():
+                summary_table.add_row(tok, f"{ratio:.2f}x")
+            console.print(summary_table)
+
+    if json_out:
+        json_out.parent.mkdir(parents=True, exist_ok=True)
+        with open(json_out, "w", encoding="utf-8") as f:
+            f.write(report.model_dump_json(indent=2))
+        console.print(f"\n[bold green]✓ Benchmark report saved to {json_out}[/bold green]")
+
+
+@app.command("route")
+def route_cmd(
+    text: Optional[str] = typer.Option(None, "--text", "-t", help="Text/prompt to route"),
+    file: Optional[Path] = typer.Option(None, "--file", "-f", help="Path to text or prompt file"),
+    budget: str = typer.Option("balanced", "--budget", "-b", help="Budget tier: frugal, balanced, frontier"),
+):
+    """Recommend the optimal model/tokenizer based on language script and syntax inflation."""
+    from token_atlas.router import recommend_model
+    if not text and not file:
+        console.print("[bold red]Error: Either --text or --file must be specified.[/bold red]")
+        raise typer.Exit(1)
+    if file:
+        text = file.read_text(encoding="utf-8")
+
+    rec = recommend_model(text, budget_tier=budget)
+    console.print(Panel(rec.explanation, title="[bold cyan]Tokenizer-Aware Model Routing Decision[/bold cyan]"))
+
+    table = Table(title="[bold]Evaluated Candidate Models & Effective Costs[/bold]")
+    table.add_column("Model ID", style="bold green")
+    table.add_column("Headline $/1M", justify="right", style="cyan")
+    table.add_column("Est. Inflation", justify="right", style="yellow")
+    table.add_column("Effective $/1M", justify="right", style="bold magenta")
+    table.add_column("Tokens Est.", justify="right", style="white")
+    table.add_column("Key Profile / Strengths", style="dim white")
+
+    for c in rec.all_candidates:
+        style = "bold green" if c.model_id == rec.recommended_model else ""
+        table.add_row(
+            f"[{style}]{c.model_id}[/{style}]" if style else c.model_id,
+            f"${c.headline_price_per_1m:.2f}",
+            f"{c.estimated_inflation:.2f}x",
+            f"${c.effective_cost_per_1m:.2f}",
+            f"{c.tokens_estimated:,}",
+            c.strengths,
+        )
+    console.print(table)
+
+
+@app.command("minify")
+def minify_cmd(
+    text: Optional[str] = typer.Option(None, "--text", "-t", help="Text/prompt to minify"),
+    file: Optional[Path] = typer.Option(None, "--file", "-f", help="Path to text or prompt file"),
+    out: Optional[Path] = typer.Option(None, "--out", "-o", help="Optional output path to write minified text"),
+    tokenizer: str = typer.Option("gpt-4o", "--tokenizer", help="Tokenizer ID to calculate token savings"),
+):
+    """Losslessly compress prompt whitespace and embedded JSON schemas to cut token tax."""
+    from token_atlas.minify import minify_prompt
+    if not text and not file:
+        console.print("[bold red]Error: Either --text or --file must be specified.[/bold red]")
+        raise typer.Exit(1)
+    if file:
+        text = file.read_text(encoding="utf-8")
+
+    res = minify_prompt(text, tokenizer_id=tokenizer)
+
+    console.print(Panel(
+        f"Original: [cyan]{res.original_char_count}[/cyan] chars ([magenta]{res.original_tokens or '—'}[/magenta] tokens)\n"
+        f"Minified: [green]{res.minified_char_count}[/green] chars ([green]{res.minified_tokens or '—'}[/green] tokens)\n"
+        f"Tokens Saved: [bold yellow]{res.tokens_saved or 0}[/bold yellow] ([bold green]{res.token_reduction_pct or 0.0}% reduction[/bold green])",
+        title="[bold cyan]Prompt Token Minification Summary[/bold cyan]"
+    ))
+
+    if out:
+        out.write_text(res.minified_text, encoding="utf-8")
+        console.print(f"[bold green]✓ Minified prompt written to {out}[/bold green]")
+    else:
+        console.print("[bold dim]Minified Output Preview:[/bold dim]")
+        console.print(res.minified_text[:500] + ("..." if len(res.minified_text) > 500 else ""))
+
+
+if __name__ == "__main__":
+    app()