Checkpoint before tooltip palette update

.claude/settings.local.json

@@ -9,7 +9,9 @@
       "Bash(git commit -m \"$\\(cat <<''EOF''\nFix Gradio compatibility for HuggingFace Spaces\n\n- Upgrade gradio to >=5.0.0 to fix API schema bug\n- Add server_name and server_port to demo.launch\\(\\)\n\nCo-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>\nEOF\n\\)\")",
       "Bash(git commit:*)",
       "Bash(git reset:*)",
-      "Bash(and top-10 predictions\" to better reflect what users see in the tooltip.\nAlso updated color legend to match the swapped model positions.\n\nCo-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>\nEOF\n\\)\")"
+      "Bash(and top-10 predictions\" to better reflect what users see in the tooltip.\nAlso updated color legend to match the swapped model positions.\n\nCo-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>\nEOF\n\\)\")",
+      "Bash(git fetch:*)",
+      "Bash(git pull:*)"
     ]
   }
 }

README.md

@@ -60,6 +60,27 @@ pip install -r requirements.txt
 python app.py
 ```
 
+## Regression Checks (Recommended)
+
+Run these after UI or rendering changes:
+
+```bash
+# Generate baseline snapshots
+conda run -n torch2 python tests/generate_snapshots.py --out tests/golden
+
+# Generate candidate snapshots
+conda run -n torch2 python tests/generate_snapshots.py --out tests/_out
+
+# Compare render-model JSON
+conda run -n torch2 python tests/compare_snapshots.py --baseline tests/golden/stress.render_model.json --candidate tests/_out/stress.render_model.json
+
+# Compare HTML output
+conda run -n torch2 python tests/compare_html.py --baseline tests/golden/stress.output.html --candidate tests/_out/stress.output.html
+
+# Optional: visual smoke placeholder
+conda run -n torch2 python tests/visual_smoke.py --html tests/_out/stress.output.html
+```
+
 ## Requirements
 
 - CUDA-capable GPU (16GB+ VRAM recommended)

app.py

@@ -7,6 +7,7 @@ Compare byte-level prediction performance between Qwen3-1.7B-Base and RWKV7-G1C-
 import gc
 import os
 from pathlib import Path
+import unicodedata
 
 import gradio as gr
 import torch
@@ -27,7 +28,7 @@ MODELS_DIR = SCRIPT_DIR / "models"
 SUPPORT_DIR = SCRIPT_DIR / "support"
 
 # Text length limits
-MAX_TEXT_LENGTH = 8192
+MAX_TEXT_LENGTH = 16384
 MIN_TEXT_LENGTH = 1
 
 # Global model cache
@@ -120,7 +121,7 @@ def validate_input(text: str) -> tuple[bool, str]:
     if not text or not text.strip():
         return False, "Please enter some text to analyze."
 
-    text = text.strip()
+    text = unicodedata.normalize("NFC", text).strip()
 
     if len(text) < MIN_TEXT_LENGTH:
         return False, f"Text is too short. Minimum {MIN_TEXT_LENGTH} characters required."
@@ -299,7 +300,15 @@ def get_default_example():
 
 
 # Build Gradio UI
-with gr.Blocks(title="Compression-Lens: RWKV-7 vs Qwen3", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(
+    title="Compression-Lens: RWKV-7 vs Qwen3",
+    theme=gr.themes.Soft(),
+    css="""
+    #input-text textarea {
+        font-family: Consolas, 'Courier New', monospace;
+    }
+    """,
+) as demo:
     gr.HTML(
         """
     <div style="text-align: center; margin-bottom: 20px;">
@@ -324,6 +333,7 @@ with gr.Blocks(title="Compression-Lens: RWKV-7 vs Qwen3", theme=gr.themes.Soft()
                 placeholder=f"Enter text to analyze (max {MAX_TEXT_LENGTH} characters)...",
                 lines=10,
                 max_lines=20,
+                elem_id="input-text",
             )
 
             with gr.Row():

core/escaping.py

@@ -0,0 +1,46 @@
+"""
+Unified escaping helpers for HTML rendering.
+
+Keep all HTML/attribute/script escaping logic in one place to avoid divergence.
+"""
+
+import json
+from typing import Any
+
+
+def escape_text_node(text: str) -> str:
+    """Escape text for HTML text node insertion."""
+    if text is None:
+        return ""
+    return (
+        text.replace("&", "&")
+        .replace("<", "&lt;")
+        .replace(">", "&gt;")
+    )
+
+
+def escape_attr(text: str) -> str:
+    """Escape text for safe placement in HTML attribute values."""
+    if text is None:
+        return ""
+    return (
+        text.replace("&", "&amp;")
+        .replace('"', "&quot;")
+        .replace("'", "&#39;")
+        .replace("<", "&lt;")
+        .replace(">", "&gt;")
+        .replace("\n", "&#10;")
+        .replace("\r", "&#13;")
+        .replace("\t", "&#9;")
+    )
+
+
+def escape_json_for_script(value: Any) -> str:
+    """Serialize JSON for safe embedding inside <script> tags."""
+    text = json.dumps(value, ensure_ascii=False)
+    # Prevent closing tags or HTML entities from breaking script context.
+    return (
+        text.replace("<", "\\u003c")
+        .replace(">", "\\u003e")
+        .replace("&", "\\u0026")
+    )

core/render_model.py

@@ -0,0 +1,59 @@
+"""
+Render model definitions for visualization.
+
+This module defines a stable, serializable intermediate representation
+between model outputs and HTML rendering.
+"""
+
+from dataclasses import dataclass, field, asdict
+from typing import Any, Dict, List, Optional
+
+
+def _has_control_chars(text: str) -> bool:
+    if not text:
+        return False
+    for ch in text:
+        code = ord(ch)
+        if code < 32 or code == 127:
+            return True
+    return False
+
+
+@dataclass
+class TokenDisplay:
+    text: str
+    kind: str  # "normal" | "control" | "raw"
+
+
+@dataclass
+class TokenInfo:
+    byte_start: int
+    byte_end: int
+    display: TokenDisplay
+    is_word: bool
+    word_id: Optional[int] = None
+    word_key: Optional[str] = None
+    bytes_hex: str = ""
+    compression: Dict[str, str] = field(default_factory=dict)
+    model_tokens: Dict[str, List[List[Any]]] = field(default_factory=dict)
+    loss: Dict[str, float] = field(default_factory=dict)
+    topk: Dict[str, Any] = field(default_factory=dict)
+    tuned_delta: float = 0.0
+
+
+@dataclass
+class RenderModel:
+    text: str
+    tokens: List[TokenInfo]
+    meta: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+
+def build_display(text: str, is_raw: bool = False) -> TokenDisplay:
+    if is_raw:
+        return TokenDisplay(text=text, kind="raw")
+    if _has_control_chars(text):
+        return TokenDisplay(text=text, kind="control")
+    return TokenDisplay(text=text, kind="normal")

core/segmentation.py

@@ -0,0 +1,24 @@
+"""
+Fallback segmentation utilities.
+
+Used for offline tests or snapshot generation when model tokenizers
+are unavailable.
+"""
+
+from typing import Dict, List
+
+
+def fallback_token_info(text: str) -> Dict[str, List]:
+    """Return minimal token info using UTF-8 codepoint boundaries."""
+    boundaries = [0]
+    byte_pos = 0
+    for ch in text:
+        byte_pos += len(ch.encode("utf-8"))
+        boundaries.append(byte_pos)
+    return {
+        "common_boundaries": boundaries,
+        "qwen_tokens": [],
+        "rwkv_tokens": [],
+        "byte_to_qwen": {},
+        "byte_to_rwkv": {},
+    }

tests/compare_html.py

@@ -0,0 +1,46 @@
+"""
+Compare two HTML files with a unified diff.
+"""
+
+import argparse
+import difflib
+from pathlib import Path
+
+
+def normalize(text: str, ignore_whitespace: bool) -> str:
+    text = text.replace("\r\n", "\n").replace("\r", "\n")
+    if ignore_whitespace:
+        return " ".join(text.split())
+    return text
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Compare HTML files.")
+    parser.add_argument("--baseline", type=Path, required=True, help="Baseline HTML path")
+    parser.add_argument("--candidate", type=Path, required=True, help="Candidate HTML path")
+    parser.add_argument("--ignore-whitespace", action="store_true", help="Normalize whitespace before diff")
+    parser.add_argument("--max-lines", type=int, default=200, help="Max diff lines to print")
+    args = parser.parse_args()
+
+    base_text = normalize(args.baseline.read_text(encoding="utf-8"), args.ignore_whitespace)
+    cand_text = normalize(args.candidate.read_text(encoding="utf-8"), args.ignore_whitespace)
+
+    base_lines = base_text.splitlines(keepends=True)
+    cand_lines = cand_text.splitlines(keepends=True)
+
+    diff = list(difflib.unified_diff(base_lines, cand_lines, fromfile=str(args.baseline), tofile=str(args.candidate)))
+
+    if not diff:
+        print("No differences found.")
+        return 0
+
+    print("Differences found:")
+    for line in diff[: args.max_lines]:
+        print(line, end="")
+    if len(diff) > args.max_lines:
+        print(f"\n... truncated ({len(diff)} total diff lines).")
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/compare_snapshots.py

@@ -0,0 +1,97 @@
+"""
+Compare two render-model JSON snapshots and report differences.
+"""
+
+import argparse
+import json
+import math
+import sys
+from pathlib import Path
+from typing import Any, List
+
+
+def load_json(path: Path) -> Any:
+    with path.open("r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def is_number(value: Any) -> bool:
+    return isinstance(value, (int, float)) and not isinstance(value, bool)
+
+
+def nearly_equal(a: float, b: float, tol: float) -> bool:
+    if math.isnan(a) and math.isnan(b):
+        return True
+    return abs(a - b) <= tol
+
+
+def compare(a: Any, b: Any, path: str, diffs: List[str], tol: float, max_diffs: int) -> None:
+    if len(diffs) >= max_diffs:
+        return
+
+    if type(a) != type(b):
+        diffs.append(f"{path}: type {type(a).__name__} != {type(b).__name__}")
+        return
+
+    if isinstance(a, dict):
+        a_keys = set(a.keys())
+        b_keys = set(b.keys())
+        for key in sorted(a_keys - b_keys):
+            diffs.append(f"{path}.{key}: missing in candidate")
+            if len(diffs) >= max_diffs:
+                return
+        for key in sorted(b_keys - a_keys):
+            diffs.append(f"{path}.{key}: extra in candidate")
+            if len(diffs) >= max_diffs:
+                return
+        for key in sorted(a_keys & b_keys):
+            compare(a[key], b[key], f"{path}.{key}", diffs, tol, max_diffs)
+            if len(diffs) >= max_diffs:
+                return
+        return
+
+    if isinstance(a, list):
+        if len(a) != len(b):
+            diffs.append(f"{path}: list length {len(a)} != {len(b)}")
+        min_len = min(len(a), len(b))
+        for idx in range(min_len):
+            compare(a[idx], b[idx], f"{path}[{idx}]", diffs, tol, max_diffs)
+            if len(diffs) >= max_diffs:
+                return
+        return
+
+    if is_number(a) and is_number(b) and tol > 0:
+        if not nearly_equal(float(a), float(b), tol):
+            diffs.append(f"{path}: {a} != {b} (tol={tol})")
+        return
+
+    if a != b:
+        diffs.append(f"{path}: {a!r} != {b!r}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Compare render-model JSON snapshots.")
+    parser.add_argument("--baseline", type=Path, required=True, help="Baseline JSON path")
+    parser.add_argument("--candidate", type=Path, required=True, help="Candidate JSON path")
+    parser.add_argument("--float-tol", type=float, default=0.0, help="Float comparison tolerance")
+    parser.add_argument("--max-diffs", type=int, default=200, help="Max diffs to display")
+    args = parser.parse_args()
+
+    baseline = load_json(args.baseline)
+    candidate = load_json(args.candidate)
+
+    diffs: List[str] = []
+    compare(baseline, candidate, "$", diffs, args.float_tol, args.max_diffs)
+
+    if diffs:
+        print(f"Differences found: {len(diffs)}")
+        for line in diffs:
+            print(line)
+        return 1
+
+    print("No differences found.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/generate_snapshots.py

@@ -0,0 +1,63 @@
+"""
+Generate offline HTML + render-model snapshots for stress inputs.
+
+Uses fallback segmentation (no model downloads).
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from core.segmentation import fallback_token_info
+from visualization.html_generator import generate_comparison_html
+SAMPLES_DIR = ROOT / "tests" / "samples"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Generate offline snapshot HTML/JSON.")
+    parser.add_argument("--out", type=Path, default=ROOT / "tests" / "golden", help="Output directory")
+    args = parser.parse_args()
+
+    input_path = SAMPLES_DIR / "stress_inputs.txt"
+    text = input_path.read_text(encoding="utf-8")
+
+    byte_len = len(text.encode("utf-8"))
+    losses_a = [0.5] * byte_len
+    losses_b = [0.6] * byte_len
+
+    html, render_model = generate_comparison_html(
+        text=text,
+        byte_losses_a=losses_a,
+        byte_losses_b=losses_b,
+        model_a_name="RWKV7 (dummy)",
+        model_b_name="Qwen3 (dummy)",
+        topk_predictions_a=None,
+        topk_predictions_b=None,
+        tokenizer_a=None,
+        tokenizer_b=None,
+        model_type_a="rwkv7",
+        model_type_b="hf",
+        token_info_override=fallback_token_info(text),
+        return_render_model=True,
+    )
+
+    out_dir = args.out
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    html_path = out_dir / "stress.output.html"
+    json_path = out_dir / "stress.render_model.json"
+
+    html_path.write_text(html, encoding="utf-8")
+    with json_path.open("w", encoding="utf-8") as f:
+        json.dump(render_model, f, ensure_ascii=False, indent=2)
+
+    print(f"Wrote {html_path}")
+    print(f"Wrote {json_path}")
+
+
+if __name__ == "__main__":
+    main()

tests/golden/README.md

@@ -0,0 +1,21 @@
+# Golden Snapshots
+
+This folder stores render-model JSON and HTML snapshots for regression checks.
+
+Suggested naming:
+- sample_01.render_model.json
+- sample_01.output.html
+
+Keep these files in sync with `tests/samples/`.
+
+Generate snapshots with:
+`conda run -n torch2 python tests/generate_snapshots.py --out tests/golden`
+
+Generate a candidate snapshot:
+`conda run -n torch2 python tests/generate_snapshots.py --out tests/_out`
+
+Compare snapshots:
+`conda run -n torch2 python tests/compare_snapshots.py --baseline tests/golden/stress.render_model.json --candidate tests/_out/stress.render_model.json`
+
+Compare HTML output:
+`conda run -n torch2 python tests/compare_html.py --baseline tests/golden/stress.output.html --candidate tests/_out/stress.output.html`

tests/samples/stress_inputs.txt

@@ -0,0 +1,27 @@
+BEGIN TEST
+  Leading spaces (2) + trailing spaces (2)··
+TAB_LITERAL: [START]	[END]   (there is a real TAB between)
+
+Raw escape-like text: \n \r \t \\n \\r \\t \\x00 \\x1f \\x7f \\xff \u0000 \u202E \u200F \u200E
+Bytes-ish hex: e5 bd 93 e7 84 b6 | 00 1f 7f ff | 0x00 0x1F 0x7F 0xFF | b"\x00\x1f\x7f\xff"
+
+HTML tags (should render as text, not tags):
+<think></think> <think>inner</think> <script>alert('x')</script> <style>body{color:red}</style>
+<div class="x" data-x="1 & 2">Hello</div> <span>Span</span> <a href="https://example.com?q=1&x=<tag>">link</a>
+<img src=x onerror=alert(1)> <br> <hr> <p>para</p> <table><tr><td>cell</td></tr></table>
+Nested-ish: </span><span data-x="</span>">confuse</span>
+
+HTML entities:
+&lt;think&gt; &lt;/think&gt; &amp; &quot; &#39; &nbsp; &#10; &#x3C; &#x3E; &#x26;
+
+Markdown-ish:
+# H1
+## H2
+- list item 1
+- list item 2
+> blockquote
+--- (three hyphens)
+
+Languages:
+中文 简体/繁體 日本語 かな カタカナ 한국어 العربية עברית हिन्दी ไทย Русский Ελληνικά Español Français Português Türkçe Việt
+RTL mix: العربية ABC עברית 123 (mixed direction)

tests/visual_smoke.py

@@ -0,0 +1,25 @@
+"""
+Placeholder for visual smoke checks.
+
+This script is intentionally minimal; hook Playwright or similar later.
+"""
+
+import argparse
+from pathlib import Path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Visual smoke check placeholder.")
+    parser.add_argument("--html", type=Path, required=True, help="HTML file to inspect")
+    args = parser.parse_args()
+
+    if not args.html.exists():
+        raise SystemExit(f"Missing HTML file: {args.html}")
+
+    print("Visual smoke placeholder:")
+    print(f"- Open in browser and visually verify: {args.html}")
+    print("- TODO: integrate Playwright for screenshot diffs.")
+
+
+if __name__ == "__main__":
+    main()

visualization/assets/main.css

@@ -0,0 +1,250 @@
+        body {
+            font-family: Consolas, 'Courier New', monospace;
+            margin: 0;
+            padding: 0;
+            background-color: #f5f5f5;
+        }
+        .header {
+            background-color: #333;
+            color: white;
+            padding: 20px;
+            position: sticky;
+            top: 0;
+            z-index: 100;
+        }
+        .header h1 {
+            margin: 0 0 15px 0;
+            font-size: 18px;
+        }
+        .meta {
+            display: flex;
+            flex-wrap: wrap;
+            gap: 20px;
+            font-size: 12px;
+            color: #c8c8c8;
+        }
+        .legend {
+            display: flex;
+            gap: 15px;
+            margin-top: 10px;
+        }
+        .legend-item {
+            display: flex;
+            align-items: center;
+            gap: 5px;
+        }
+        .legend-box {
+            width: 20px;
+            height: 12px;
+            border: 1px solid #666;
+        }
+        .content {
+            background-color: white;
+            margin: 10px;
+            padding: 15px;
+            border: 1px solid #ccc;
+            font-size: 14px;
+            line-height: 1.8;
+            white-space: pre-wrap;
+            overflow-wrap: anywhere;
+            word-wrap: break-word;
+            position: relative;
+        }
+        .content span {
+            padding: 1px 0;
+        }
+        .word {
+            cursor: pointer;
+            position: relative;
+        }
+        .word:hover {
+            outline: 2px solid #007bff;
+            outline-offset: 1px;
+        }
+        .word.highlighted {
+            outline: 2px solid #ff6b6b;
+            outline-offset: 1px;
+        }
+        #svg-overlay {
+            position: fixed;
+            top: 0;
+            left: 0;
+            width: 100%;
+            height: 100%;
+            pointer-events: none;
+            z-index: 1000;
+        }
+        .link-line {
+            stroke: #007bff;
+            stroke-width: 2;
+            fill: none;
+            opacity: 0.7;
+        }
+        .link-dot {
+            fill: #007bff;
+            opacity: 0.8;
+        }
+        .token {
+            position: relative;
+            cursor: help;
+        }
+        .token:hover {
+            outline: 1px dashed #666;
+        }
+        .token-kind-control {
+            color: #f59e0b;
+        }
+        .token-kind-raw {
+            color: #fb7185;
+        }
+        #tooltip {
+            position: fixed;
+            background-color: rgba(0, 0, 0, 0.9);
+            color: white;
+            padding: 10px 14px;
+            border-radius: 6px;
+            font-size: 12px;
+            max-width: none;
+            width: max-content;
+            z-index: 2000;
+            pointer-events: none;
+            display: none;
+            line-height: 1.6;
+            box-shadow: 0 2px 10px rgba(0,0,0,0.3);
+        }
+        #tooltip .label {
+            color: #aaa;
+            font-weight: bold;
+        }
+        #tooltip .bytes {
+            color: #a5f3fc;
+            font-family: monospace;
+        }
+        #tooltip .loss-a {
+            color: #86efac;
+            font-family: monospace;
+        }
+        #tooltip .loss-b {
+            color: #fca5a5;
+            font-family: monospace;
+        }
+        #tooltip .model-a {
+            color: #fcd34d;
+        }
+        #tooltip .model-b {
+            color: #7dd3fc;
+        }
+        #tooltip .topk-section {
+            margin-top: 8px;
+            padding-top: 8px;
+            border-top: 1px solid #555;
+        }
+        #tooltip .topk-container {
+            display: flex;
+            gap: 16px;
+        }
+        #tooltip .topk-column {
+            flex: 1;
+            min-width: 180px;
+        }
+        #tooltip .topk-title {
+            color: #aaa;
+            font-weight: bold;
+            margin-bottom: 4px;
+            font-size: 11px;
+        }
+        #tooltip .topk-title.model-a {
+            color: #86efac;
+        }
+        #tooltip .topk-title.model-b {
+            color: #fca5a5;
+        }
+        #tooltip .topk-list {
+            font-size: 11px;
+        }
+        #tooltip .topk-item {
+            display: flex;
+            gap: 4px;
+            padding: 1px 0;
+            align-items: flex-start;
+        }
+        #tooltip .token-block {
+            margin-top: 6px;
+            display: flex;
+            align-items: center;
+            gap: 6px;
+            white-space: nowrap;
+            flex-wrap: nowrap;
+            overflow-x: visible;
+        }
+        #tooltip .token-chips {
+            display: flex;
+            flex-wrap: nowrap;
+            gap: 4px;
+            align-items: center;
+            flex: 0 0 auto;
+        }
+        #tooltip .token-chip-group {
+            display: inline-flex;
+            align-items: center;
+            gap: 4px;
+            flex: 0 0 auto;
+            white-space: nowrap;
+        }
+        #tooltip .token-prob {
+            color: #86efac;
+            font-family: monospace;
+            font-size: 11px;
+            white-space: nowrap;
+        }
+        #tooltip .token-id {
+            color: #888;
+            font-family: monospace;
+            white-space: nowrap;
+        }
+        #tooltip .token-chip {
+            max-width: 100%;
+        }
+        #tooltip .token-chip-group .topk-token {
+            white-space: pre;
+            overflow-wrap: normal;
+            word-break: normal;
+        }
+        #tooltip .topk-rank {
+            color: #888;
+            min-width: 18px;
+        }
+        #tooltip .topk-rank.hit {
+            color: #ffd700;
+        }
+        #tooltip .topk-token {
+            color: #a5f3fc;
+            white-space: pre-wrap;
+            overflow-wrap: anywhere;
+            word-break: break-word;
+            font-family: monospace;
+            background-color: rgba(255, 255, 255, 0.08);
+            padding: 0 4px;
+            border-radius: 3px;
+            display: inline-block;
+            max-width: 100%;
+        }
+        #tooltip .esc-control {
+            color: #fbbf24;
+        }
+        #tooltip .esc-raw {
+            color: #fb7185;
+        }
+        #tooltip .topk-prob {
+            color: #86efac;
+            min-width: 45px;
+            text-align: right;
+        }
+        #tooltip .topk-hit {
+            color: #22c55e;
+        }
+        #tooltip .topk-miss {
+            color: #ef4444;
+            font-style: italic;
+        }
+    

visualization/assets/main.js

@@ -0,0 +1,615 @@
+        const svgOverlay = document.getElementById('svg-overlay');
+        const content = document.querySelector('.content');
+        const renderModelEl = document.getElementById('render-model');
+        let renderModel = {};
+        let renderTokens = [];
+        try {
+            renderModel = JSON.parse(renderModelEl ? renderModelEl.textContent : '{}') || {};
+            renderTokens = Array.isArray(renderModel.tokens) ? renderModel.tokens : [];
+        } catch (e) {
+            console.warn('Failed to parse render model JSON:', e);
+            renderModel = {};
+            renderTokens = [];
+        }
+
+        function escapeControlChars(text) {
+            if (!text) return text;
+            let out = '';
+            for (let i = 0; i < text.length; i++) {
+                const ch = text[i];
+                const code = text.charCodeAt(i);
+                if (ch === '\\') {
+                    out += '\\\\';
+                } else if (ch === '\n') {
+                    out += '\\n';
+                } else if (ch === '\r') {
+                    out += '\\r';
+                } else if (ch === '\t') {
+                    out += '\\t';
+                } else if (code < 32 || code === 127) {
+                    out += '\\x' + code.toString(16).padStart(2, '0');
+                } else {
+                    out += ch;
+                }
+            }
+            return out;
+        }
+
+        if (content) {
+            while (content.firstChild) {
+                content.removeChild(content.firstChild);
+            }
+            const wordCounts = {};
+            renderTokens.forEach((token) => {
+                if (token && token.is_word && token.word_key) {
+                    wordCounts[token.word_key] = (wordCounts[token.word_key] || 0) + 1;
+                }
+            });
+
+            renderTokens.forEach((token, idx) => {
+                const span = document.createElement('span');
+                span.className = 'token';
+                span.dataset.tokenIdx = String(idx);
+                span.dataset.tunedDelta = (token && typeof token.tuned_delta === 'number') ? String(token.tuned_delta) : '0';
+                const kind = (token && token.display && token.display.kind) ? token.display.kind : 'normal';
+                const text = (token && token.display && typeof token.display.text === 'string') ? token.display.text : '';
+                const hasVisible = (() => {
+                    if (!text) return false;
+                    for (let i = 0; i < text.length; i++) {
+                        const code = text.charCodeAt(i);
+                        if (code >= 32 && code !== 127) {
+                            return true;
+                        }
+                    }
+                    return false;
+                })();
+                const mainKind = (kind === 'control' && hasVisible) ? 'normal' : kind;
+                if (text.includes('\n') || text.includes('\r')) {
+                    span.dataset.hasLinebreak = '1';
+                }
+                if (mainKind === 'control') {
+                    span.classList.add('token-kind-control');
+                    span.textContent = text;
+                } else if (mainKind === 'raw') {
+                    span.classList.add('token-kind-raw');
+                    span.textContent = text;
+                } else {
+                    span.textContent = text;
+                }
+
+                if (token && token.is_word && token.word_key && wordCounts[token.word_key] > 1) {
+                    span.classList.add('word');
+                    span.dataset.word = token.word_key;
+                    if (token.word_id !== undefined && token.word_id !== null) {
+                        span.dataset.wordId = String(token.word_id);
+                    }
+                }
+
+                content.appendChild(span);
+            });
+        }
+
+        const words = document.querySelectorAll('.word');
+        const wordGroups = {};
+        words.forEach(word => {
+            const wordText = word.getAttribute('data-word');
+            if (!wordGroups[wordText]) {
+                wordGroups[wordText] = [];
+            }
+            wordGroups[wordText].push(word);
+        });
+
+        function clearLines() {
+            while (svgOverlay.firstChild) {
+                svgOverlay.removeChild(svgOverlay.firstChild);
+            }
+            words.forEach(w => w.classList.remove('highlighted'));
+        }
+
+        function pickRectByY(rects, targetY) {
+            if (!rects || rects.length === 0) return null;
+            let best = rects[0];
+            let bestDist = Infinity;
+            rects.forEach(r => {
+                const cy = r.top + r.height / 2;
+                const dist = Math.abs(cy - targetY);
+                if (dist < bestDist) {
+                    best = r;
+                    bestDist = dist;
+                }
+            });
+            return best;
+        }
+
+        function getAnchorRect(element, targetY) {
+            const rects = Array.from(element.getClientRects());
+            if (rects.length === 0) return element.getBoundingClientRect();
+            if (rects.length === 1) return rects[0];
+            const picked = pickRectByY(rects, targetY);
+            return picked || rects[0];
+        }
+
+        function drawLines(hoveredWord, evt) {
+            clearLines();
+
+            const wordText = hoveredWord.getAttribute('data-word');
+            const wordId = parseInt(hoveredWord.getAttribute('data-word-id'));
+            const sameWords = wordGroups[wordText] || [];
+
+            const previousWords = sameWords.filter(w => {
+                const id = parseInt(w.getAttribute('data-word-id'));
+                return id < wordId;
+            });
+
+            if (previousWords.length === 0) return;
+
+            sameWords.forEach(w => w.classList.add('highlighted'));
+
+            const targetY = evt ? evt.clientY : (hoveredWord.getBoundingClientRect().top + hoveredWord.getBoundingClientRect().height / 2);
+            const hoveredRect = getAnchorRect(hoveredWord, targetY);
+            const hoveredX = hoveredRect.left + hoveredRect.width / 2;
+            const hoveredY = hoveredRect.top + hoveredRect.height / 2;
+
+            previousWords.forEach(prevWord => {
+                const prevRect = getAnchorRect(prevWord, hoveredY);
+                const prevX = prevRect.left + prevRect.width / 2;
+                const prevY = prevRect.top + prevRect.height / 2;
+
+                const midX = (hoveredX + prevX) / 2;
+                const midY = Math.min(hoveredY, prevY) - 30;
+
+                const path = document.createElementNS('http://www.w3.org/2000/svg', 'path');
+                path.setAttribute('class', 'link-line');
+                path.setAttribute('d', `M ${prevX} ${prevY} Q ${midX} ${midY} ${hoveredX} ${hoveredY}`);
+                svgOverlay.appendChild(path);
+
+                const dot1 = document.createElementNS('http://www.w3.org/2000/svg', 'circle');
+                dot1.setAttribute('class', 'link-dot');
+                dot1.setAttribute('cx', prevX);
+                dot1.setAttribute('cy', prevY);
+                dot1.setAttribute('r', 4);
+                svgOverlay.appendChild(dot1);
+
+                const dot2 = document.createElementNS('http://www.w3.org/2000/svg', 'circle');
+                dot2.setAttribute('class', 'link-dot');
+                dot2.setAttribute('cx', hoveredX);
+                dot2.setAttribute('cy', hoveredY);
+                dot2.setAttribute('r', 4);
+                svgOverlay.appendChild(dot2);
+            });
+        }
+
+        words.forEach(word => {
+            word.addEventListener('mouseenter', (e) => drawLines(word, e));
+            word.addEventListener('mouseleave', clearLines);
+        });
+
+        window.addEventListener('scroll', clearLines);
+
+        const tooltip = document.getElementById('tooltip');
+        const tokenSpans = document.querySelectorAll('.token');
+
+        tokenSpans.forEach(token => {
+            token.addEventListener('mouseenter', (e) => {
+                const tokenIdx = parseInt(token.dataset.tokenIdx);
+                const tokenInfo = (!isNaN(tokenIdx) && renderTokens[tokenIdx]) ? renderTokens[tokenIdx] : null;
+                const bytes = (tokenInfo && tokenInfo.bytes_hex) ? tokenInfo.bytes_hex : '';
+                const compressionA = (tokenInfo && tokenInfo.compression && tokenInfo.compression.rwkv) ? tokenInfo.compression.rwkv : '';
+                const compressionB = (tokenInfo && tokenInfo.compression && tokenInfo.compression.qwen) ? tokenInfo.compression.qwen : '';
+                const avgCompressionA = (tokenInfo && tokenInfo.loss && typeof tokenInfo.loss.rwkv === 'number') ? tokenInfo.loss.rwkv.toFixed(2) : '';
+                const avgCompressionB = (tokenInfo && tokenInfo.loss && typeof tokenInfo.loss.qwen === 'number') ? tokenInfo.loss.qwen.toFixed(2) : '';
+                const modelA = (tokenInfo && tokenInfo.model_tokens && tokenInfo.model_tokens.rwkv) ? tokenInfo.model_tokens.rwkv : null;
+                const modelB = (tokenInfo && tokenInfo.model_tokens && tokenInfo.model_tokens.qwen) ? tokenInfo.model_tokens.qwen : null;
+                const top5A = (tokenInfo && tokenInfo.topk && tokenInfo.topk.rwkv) ? tokenInfo.topk.rwkv : null;
+                const top5B = (tokenInfo && tokenInfo.topk && tokenInfo.topk.qwen) ? tokenInfo.topk.qwen : null;
+
+                function hasControlChars(text) {
+                    if (!text) return false;
+                    for (let i = 0; i < text.length; i++) {
+                        const code = text.charCodeAt(i);
+                        if (code < 32 || code === 127) {
+                            return true;
+                        }
+                    }
+                    return false;
+                }
+
+                function resolveKind(text, kindHint) {
+                    if (kindHint === 'raw' || kindHint === 'control' || kindHint === 'normal') {
+                        return kindHint;
+                    }
+                    if (kindHint === true) {
+                        return 'raw';
+                    }
+                    if (hasControlChars(text)) {
+                        return 'control';
+                    }
+                    return 'normal';
+                }
+
+                function appendEscapedWithControlColor(container, text) {
+                    if (text === undefined || text === null) return;
+                    let buffer = '';
+                    const flush = () => {
+                        if (buffer) {
+                            container.appendChild(document.createTextNode(buffer));
+                            buffer = '';
+                        }
+                    };
+                    for (let i = 0; i < text.length; i++) {
+                        const ch = text[i];
+                        const code = text.charCodeAt(i);
+                        if (ch === '\\') {
+                            buffer += '\\\\';
+                            continue;
+                        }
+                        if (ch === '\n' || ch === '\r' || ch === '\t' || code < 32 || code === 127) {
+                            flush();
+                            const span = document.createElement('span');
+                            span.className = 'esc-control';
+                            if (ch === '\n') {
+                                span.textContent = '\\n';
+                            } else if (ch === '\r') {
+                                span.textContent = '\\r';
+                            } else if (ch === '\t') {
+                                span.textContent = '\\t';
+                            } else {
+                                span.textContent = '\\x' + code.toString(16).padStart(2, '0');
+                            }
+                            container.appendChild(span);
+                            continue;
+                        }
+                        buffer += ch;
+                    }
+                    flush();
+                }
+
+                function appendTokenText(container, text, kindHint) {
+                    const display = (text !== undefined && text !== null) ? text : '';
+                    const kind = resolveKind(display, kindHint);
+                    while (container.firstChild) {
+                        container.removeChild(container.firstChild);
+                    }
+                    if (kind === 'raw') {
+                        const span = document.createElement('span');
+                        span.className = 'esc-raw';
+                        span.textContent = display;
+                        container.appendChild(span);
+                        return;
+                    }
+                    if (kind === 'control') {
+                        appendEscapedWithControlColor(container, display);
+                        return;
+                    }
+                    container.textContent = display;
+                }
+
+                function formatTopkColumn(topkData, modelName, titleClass) {
+                    const column = document.createElement('div');
+                    column.className = 'topk-column';
+                    const title = document.createElement('div');
+                    title.className = 'topk-title ' + titleClass;
+                    title.textContent = modelName;
+                    column.appendChild(title);
+                    const list = document.createElement('div');
+                    list.className = 'topk-list';
+                    column.appendChild(list);
+
+                    if (!topkData) {
+                        list.textContent = 'N/A';
+                        return column;
+                    }
+                    try {
+                        const data = topkData;
+                        let actualId = null;
+                        let rank = null;
+                        let actualProb = null;
+                        let topkList = [];
+                        if (data.length >= 4) {
+                            [actualId, rank, actualProb, topkList] = data;
+                        } else {
+                            [actualId, rank, topkList] = data;
+                        }
+                        topkList.forEach((item, idx) => {
+                            const tokenId = item[0];
+                            const prob = item[1];
+                            const tokenText = item[2];
+                            const isRaw = item.length > 3 ? item[3] : false;
+                            const isHit = tokenId === actualId;
+                            const rankClass = isHit ? 'topk-rank hit' : 'topk-rank';
+                            const rawText = (tokenText !== undefined && tokenText !== null) ? tokenText : '';
+                            const displayText = (rawText !== '') ? rawText : ('[' + tokenId + ']');
+
+                            const row = document.createElement('div');
+                            row.className = 'topk-item';
+
+                            const rankSpan = document.createElement('span');
+                            rankSpan.className = rankClass;
+                            rankSpan.textContent = (idx + 1) + '.';
+                            row.appendChild(rankSpan);
+
+                            const tokenSpan = document.createElement('span');
+                            tokenSpan.className = 'topk-token';
+                            tokenSpan.title = 'ID: ' + tokenId;
+                            appendTokenText(tokenSpan, displayText, isRaw);
+                            row.appendChild(tokenSpan);
+
+                            const probSpan = document.createElement('span');
+                            probSpan.className = 'topk-prob';
+                            probSpan.textContent = (prob * 100).toFixed(2) + '%';
+                            row.appendChild(probSpan);
+
+                            if (isHit) {
+                                const hit = document.createElement('span');
+                                hit.className = 'topk-hit';
+                                hit.textContent = '✓';
+                                row.appendChild(hit);
+                            }
+
+                            list.appendChild(row);
+                        });
+
+                        if (rank > 10) {
+                            let probSuffix = '';
+                            const probVal = parseFloat(actualProb);
+                            if (!isNaN(probVal)) {
+                                probSuffix = ' (' + (probVal * 100).toFixed(4) + '%)';
+                            }
+                            const miss = document.createElement('div');
+                            miss.className = 'topk-item topk-miss';
+                            miss.textContent = 'Actual rank: ' + rank + probSuffix;
+                            list.appendChild(miss);
+                        }
+                        return column;
+                    } catch (e) {
+                        console.error('Error in formatTopkColumn for ' + modelName + ':', e);
+                        console.error('topkData:', topkData);
+                        list.textContent = 'Error: ' + e.message;
+                        return column;
+                    }
+                }
+
+                function formatTokenChips(modelData, label, labelClass) {
+                    const block = document.createElement('div');
+                    block.className = 'token-block';
+                    const labelSpan = document.createElement('span');
+                    labelSpan.className = 'label ' + labelClass;
+                    labelSpan.textContent = label + ':';
+                    block.appendChild(labelSpan);
+
+                    const chips = document.createElement('div');
+                    chips.className = 'token-chips';
+                    block.appendChild(chips);
+
+                    if (!modelData) {
+                        const na = document.createElement('span');
+                        na.className = 'topk-token token-chip';
+                        na.textContent = 'N/A';
+                        chips.appendChild(na);
+                        return block;
+                    }
+                    try {
+                        const tokenList = modelData;
+                        tokenList.forEach((item) => {
+                            const tokenId = item[0];
+                            const tokenText = item[1];
+                            const kindHint = item.length > 2 ? item[2] : false;
+                            const probVal = item.length > 3 ? item[3] : null;
+                            const displayText = (tokenText !== undefined && tokenText !== null) ? tokenText : '';
+
+                            const group = document.createElement('span');
+                            group.className = 'token-chip-group';
+                            group.title = 'ID: ' + tokenId;
+
+                            const idSpan = document.createElement('span');
+                            idSpan.className = 'token-id';
+                            idSpan.textContent = '[' + tokenId + ']';
+                            group.appendChild(idSpan);
+
+                            const chipSpan = document.createElement('span');
+                            chipSpan.className = 'topk-token token-chip';
+                            appendTokenText(chipSpan, displayText, kindHint);
+                            group.appendChild(chipSpan);
+
+                            if (probVal !== null && probVal !== undefined) {
+                                const probSpan = document.createElement('span');
+                                probSpan.className = 'token-prob';
+                                const numericProb = typeof probVal === 'number' ? probVal : parseFloat(probVal);
+                                if (!isNaN(numericProb)) {
+                                    probSpan.textContent = (numericProb * 100).toFixed(2) + '%';
+                                } else {
+                                    probSpan.textContent = String(probVal);
+                                }
+                                group.appendChild(probSpan);
+                            }
+
+                            chips.appendChild(group);
+                        });
+                        return block;
+                    } catch (e) {
+                        console.error('Error in formatTokenChips for ' + label + ':', e);
+                        console.error('modelData:', modelData);
+                        const err = document.createElement('span');
+                        err.className = 'topk-token token-chip';
+                        err.textContent = 'Error: ' + e.message;
+                        chips.appendChild(err);
+                        return block;
+                    }
+                }
+
+                tooltip.replaceChildren();
+
+                const bytesRow = document.createElement('div');
+                const bytesLabel = document.createElement('span');
+                bytesLabel.className = 'label';
+                bytesLabel.textContent = 'Bytes:';
+                const bytesValue = document.createElement('span');
+                bytesValue.className = 'bytes';
+                bytesValue.textContent = bytes || '(empty)';
+                bytesRow.appendChild(bytesLabel);
+                bytesRow.appendChild(document.createTextNode(' '));
+                bytesRow.appendChild(bytesValue);
+                tooltip.appendChild(bytesRow);
+
+                const rwkvRow = document.createElement('div');
+                const rwkvLabel = document.createElement('span');
+                rwkvLabel.className = 'label';
+                rwkvLabel.textContent = 'RWKV Compression Rate:';
+                const rwkvValue = document.createElement('span');
+                rwkvValue.className = 'loss-a';
+                rwkvValue.textContent = compressionA || '(empty)';
+                if (avgCompressionA) {
+                    rwkvValue.textContent += ' (avg: ' + avgCompressionA + '%)';
+                }
+                rwkvRow.appendChild(rwkvLabel);
+                rwkvRow.appendChild(document.createTextNode(' '));
+                rwkvRow.appendChild(rwkvValue);
+                tooltip.appendChild(rwkvRow);
+
+                const qwenRow = document.createElement('div');
+                const qwenLabel = document.createElement('span');
+                qwenLabel.className = 'label';
+                qwenLabel.textContent = 'Qwen Compression Rate:';
+                const qwenValue = document.createElement('span');
+                qwenValue.className = 'loss-b';
+                qwenValue.textContent = compressionB || '(empty)';
+                if (avgCompressionB) {
+                    qwenValue.textContent += ' (avg: ' + avgCompressionB + '%)';
+                }
+                qwenRow.appendChild(qwenLabel);
+                qwenRow.appendChild(document.createTextNode(' '));
+                qwenRow.appendChild(qwenValue);
+                tooltip.appendChild(qwenRow);
+
+                const hr = document.createElement('hr');
+                hr.style.borderColor = '#555';
+                hr.style.margin = '6px 0';
+                tooltip.appendChild(hr);
+
+                tooltip.appendChild(formatTokenChips(modelA, 'RWKV', 'model-a'));
+                tooltip.appendChild(formatTokenChips(modelB, 'Qwen', 'model-b'));
+
+                if (top5A || top5B) {
+                    const topkSection = document.createElement('div');
+                    topkSection.className = 'topk-section';
+                    const topkContainer = document.createElement('div');
+                    topkContainer.className = 'topk-container';
+                    topkContainer.appendChild(formatTopkColumn(top5A, 'RWKV Top10', 'model-a'));
+                    topkContainer.appendChild(formatTopkColumn(top5B, 'Qwen Top10', 'model-b'));
+                    topkSection.appendChild(topkContainer);
+                    tooltip.appendChild(topkSection);
+                }
+
+                tooltip.style.display = 'block';
+            });
+
+            token.addEventListener('mousemove', (e) => {
+                const tooltipRect = tooltip.getBoundingClientRect();
+                const viewportWidth = window.innerWidth;
+                const viewportHeight = window.innerHeight;
+
+                let x = e.clientX + 15;
+                let y = e.clientY + 15;
+
+                if (x + tooltipRect.width > viewportWidth - 10) {
+                    x = e.clientX - tooltipRect.width - 15;
+                }
+                if (y + tooltipRect.height > viewportHeight - 10) {
+                    y = e.clientY - tooltipRect.height - 15;
+                }
+                if (x < 10) x = 10;
+                if (y < 10) y = 10;
+
+                tooltip.style.left = x + 'px';
+                tooltip.style.top = y + 'px';
+            });
+
+            token.addEventListener('mouseleave', () => {
+                tooltip.style.display = 'none';
+            });
+        });
+
+        const slider = document.getElementById('color-range-slider');
+        const rangeValue = document.getElementById('color-range-value');
+
+        // Collect all tuned_delta values
+        const tokenData = [];
+        const linebreakTokens = [];
+        tokenSpans.forEach((token, idx) => {
+            if (token.dataset.hasLinebreak === '1') {
+                linebreakTokens.push(token);
+                return;
+            }
+            const tunedDelta = parseFloat(token.dataset.tunedDelta);
+            if (!isNaN(tunedDelta)) {
+                tokenData.push({ token, tunedDelta, absDelta: Math.abs(tunedDelta) });
+            }
+        });
+
+        // Calculate max_abs_tuned_delta for normalization
+        const maxAbsDelta = Math.max(...tokenData.map(d => d.absDelta), 1e-9);
+
+        // Sort by |tuned_delta| to get rankings
+        const sortedByAbs = [...tokenData].sort((a, b) => b.absDelta - a.absDelta);
+        sortedByAbs.forEach((item, rank) => {
+            item.rank = rank;  // rank 0 = largest deviation
+        });
+
+        function tunedDeltaToColor(tunedDelta, maxAbsDelta, exponent) {
+            // Normalize to [-1, 1]
+            const normalized = Math.max(-1, Math.min(1, tunedDelta / maxAbsDelta));
+            let r, g, b;
+            if (normalized < 0) {
+                // Green (RWKV better)
+                const intensity = Math.pow(-normalized, exponent);
+                r = Math.round(255 * (1 - intensity * 0.85));
+                g = 255;
+                b = Math.round(255 * (1 - intensity * 0.85));
+            } else {
+                // Red (RWKV worse)
+                const intensity = Math.pow(normalized, exponent);
+                r = 255;
+                g = Math.round(255 * (1 - intensity * 0.85));
+                b = Math.round(255 * (1 - intensity * 0.85));
+            }
+            return `rgb(${r}, ${g}, ${b})`;
+        }
+
+        function updateColors(colorRangePercent) {
+            // colorRangePercent: 0-100, represents the proportion of tokens to color
+            const colorCount = Math.round(tokenData.length * colorRangePercent / 100);
+
+            // Calculate exponent: 100% -> 0.5, 0% -> 1.0
+            const exponent = 1 - (colorRangePercent / 100) * 0.5;
+
+            // Calculate max deviation within the colored range
+            let maxAbsDeltaInRange = 1e-9;
+            tokenData.forEach(item => {
+                if (item.rank < colorCount) {
+                    maxAbsDeltaInRange = Math.max(maxAbsDeltaInRange, item.absDelta);
+                }
+            });
+
+            tokenData.forEach(item => {
+                if (item.rank < colorCount) {
+                    // Use dynamic normalization based on colored range
+                    item.token.style.backgroundColor = tunedDeltaToColor(item.tunedDelta, maxAbsDeltaInRange, exponent);
+                } else {
+                    // Outside color range, white
+                    item.token.style.backgroundColor = 'rgb(255, 255, 255)';
+                }
+            });
+            linebreakTokens.forEach(token => {
+                token.style.backgroundColor = 'rgb(255, 255, 255)';
+            });
+        }
+
+        slider.addEventListener('input', (e) => {
+            const val = parseFloat(e.target.value);
+            rangeValue.textContent = val.toFixed(1) + '%';
+            updateColors(val);
+        });
+
+        // Apply default color range on page load
+        updateColors(10);
+    

visualization/html_generator.py

@@ -4,16 +4,22 @@ HTML visualization generator for UncheatableEval.
 Generates interactive HTML visualizations comparing byte-level losses between two models.
 """
 
-import base64
+import bisect
 import json
 import math
 import re
+from pathlib import Path
 from typing import List, Tuple, Optional, Set
 
 import numpy as np
 
+from core.escaping import escape_json_for_script
+from core.render_model import RenderModel, TokenInfo, build_display
+from visualization.render import render_page
 from core.helpers import TokenizerBytesConverter
 
+ASSETS_DIR = Path(__file__).resolve().parent / "assets"
+
 
 # Compression rate conversion factor
 COMPRESSION_RATE_FACTOR = (1.0 / math.log(2.0)) * 0.125 * 100.0
@@ -113,10 +119,17 @@ def get_token_info_for_text(text: str) -> dict:
     qwen_boundaries = set([0] + [t[1] for t in qwen_tokens])
     rwkv_boundaries = set([0] + [t[1] for t in rwkv_tokens])
     utf8_boundaries = set([0])
+    whitespace_boundaries = set()
+    linebreak_boundaries = set()
     byte_pos = 0
     for ch in text:
-        byte_pos += len(ch.encode("utf-8"))
+        ch_bytes = ch.encode("utf-8")
+        byte_pos += len(ch_bytes)
         utf8_boundaries.add(byte_pos)
+        if ch.isspace():
+            whitespace_boundaries.add(byte_pos)
+        if ch in ("\n", "\r"):
+            linebreak_boundaries.add(byte_pos)
     common_boundaries = sorted(qwen_boundaries & rwkv_boundaries & utf8_boundaries)
     # Ensure we always include the end boundary
     text_end = len(text.encode("utf-8"))
@@ -124,6 +137,75 @@ def get_token_info_for_text(text: str) -> dict:
         common_boundaries.append(text_end)
         common_boundaries = sorted(common_boundaries)
 
+    # Refine overly large segments to avoid giant spans in the UI.
+    max_segment_bytes = 24
+    utf8_sorted = sorted(utf8_boundaries)
+    linebreak_sorted = sorted(linebreak_boundaries)
+
+    def split_by_max(start: int, end: int) -> List[int]:
+        if end - start <= max_segment_bytes:
+            return [end]
+        left = bisect.bisect_right(utf8_sorted, start)
+        right = bisect.bisect_left(utf8_sorted, end)
+        candidates = utf8_sorted[left:right]
+        if not candidates:
+            return [end]
+        out = []
+        pos = start
+        idx = 0
+        while pos < end:
+            limit = min(end, pos + max_segment_bytes)
+            j = bisect.bisect_right(candidates, limit) - 1
+            if j < idx:
+                out.append(end)
+                break
+            split_at = None
+            for k in range(j, idx - 1, -1):
+                if candidates[k] in whitespace_boundaries:
+                    split_at = candidates[k]
+                    j = k
+                    break
+            if split_at is None:
+                split_at = candidates[j]
+            if split_at <= pos:
+                split_at = candidates[j]
+            out.append(split_at)
+            pos = split_at
+            idx = j + 1
+            if pos >= end:
+                break
+            if idx >= len(candidates):
+                out.append(end)
+                break
+        if not out:
+            out = [end]
+        elif out[-1] != end:
+            out.append(end)
+        return out
+
+    def split_segment(start: int, end: int) -> List[int]:
+        if start >= end:
+            return []
+        lb_left = bisect.bisect_right(linebreak_sorted, start)
+        lb_right = bisect.bisect_left(linebreak_sorted, end)
+        linebreaks = linebreak_sorted[lb_left:lb_right]
+        if not linebreaks:
+            return split_by_max(start, end)
+        out = []
+        seg_start = start
+        for lb in linebreaks:
+            out.extend(split_by_max(seg_start, lb))
+            seg_start = lb
+        out.extend(split_by_max(seg_start, end))
+        return out
+
+    refined_boundaries = [common_boundaries[0]] if common_boundaries else [0]
+    for i in range(len(common_boundaries) - 1):
+        start = common_boundaries[i]
+        end = common_boundaries[i + 1]
+        refined_boundaries.extend(split_segment(start, end))
+    common_boundaries = sorted(set(refined_boundaries))
+
     return {
         "common_boundaries": common_boundaries,
         "qwen_tokens": qwen_tokens,
@@ -145,6 +227,8 @@ def generate_comparison_html(
     tokenizer_b=None,
     model_type_a: str = "hf",
     model_type_b: str = "rwkv7",
+    token_info_override: Optional[dict] = None,
+    return_render_model: bool = False,
 ) -> str:
     """
     Generate an interactive HTML visualization comparing two models.
@@ -161,9 +245,11 @@ def generate_comparison_html(
         tokenizer_b: Tokenizer for model B
         model_type_a: Type of model A ("hf" or "rwkv7")
         model_type_b: Type of model B ("hf" or "rwkv7")
+        token_info_override: Optional precomputed token info (for offline tests).
+        return_render_model: If True, return (html, render_model_dict)
 
     Returns:
-        HTML string with interactive visualization
+        HTML string with interactive visualization, or (html, render_model_dict) if return_render_model=True
     """
 
     def decode_token(token_id: int, tokenizer, model_type: str) -> Tuple[str, bool]:
@@ -197,7 +283,12 @@ def generate_comparison_html(
         try:
             if model_type in ["rwkv", "rwkv7"]:
                 # RWKV tokenizer provides raw bytes
-                token_bytes = tokenizer.decodeBytes([token_id])
+                try:
+                    token_bytes = tokenizer.decodeBytes([token_id])
+                except Exception as e:
+                    if token_id == 0:
+                        return f"[{token_id}]", False
+                    raise e
                 if token_bytes:
                     try:
                         decoded = token_bytes.decode("utf-8")
@@ -290,7 +381,7 @@ def generate_comparison_html(
 
     # Get token info
     text_bytes = text.encode("utf-8")
-    token_info = get_token_info_for_text(text)
+    token_info = token_info_override if token_info_override is not None else get_token_info_for_text(text)
     common_boundaries = token_info["common_boundaries"]
     qwen_tokens = token_info["qwen_tokens"]
     rwkv_tokens = token_info["rwkv_tokens"]
@@ -301,14 +392,13 @@ def generate_comparison_html(
 
     def get_tokens_for_range(byte_start, byte_end, token_list):
         result = []
-        for t_start, t_end, token_id, t_bytes in token_list:
+        for idx, (t_start, t_end, token_id, t_bytes) in enumerate(token_list):
             if t_start < byte_end and t_end > byte_start:
-                result.append((token_id, t_bytes))
+                result.append((idx, token_id, t_bytes))
         return result
 
     # Build tokens based on common boundaries
     tokens = []
-    token_count = 0
     for i in range(len(common_boundaries) - 1):
         start_byte = common_boundaries[i]
         end_byte = common_boundaries[i + 1]
@@ -361,22 +451,8 @@ def generate_comparison_html(
             token["word_id"] = word_id_counter
             word_id_counter += 1
 
-    # Build HTML content
-    html_content = []
-
-    def escape_for_attr(s):
-        # Escape all characters that could break HTML attributes
-        # Order matters: & must be escaped first
-        return (
-            s.replace("&", "&amp;")
-            .replace('"', "&quot;")
-            .replace("'", "&#39;")
-            .replace("<", "&lt;")
-            .replace(">", "&gt;")
-            .replace("\n", "&#10;")
-            .replace("\r", "&#13;")
-            .replace("\t", "&#9;")
-        )
+    # Build render model (HTML content built in JS)
+    render_tokens = []
 
     for token in tokens:
         token_text = token["text"]
@@ -402,18 +478,6 @@ def generate_comparison_html(
             except UnicodeDecodeError:
                 return "".join([f"\\x{b:02x}" for b in token_bytes]), True
 
-        # Model A (RWKV7) - tokens overlapping this byte range
-        model_a_info = ""
-        if token["rwkv_tokens"]:
-            model_a_list = [[tid, *token_bytes_to_display_text(tb)] for tid, tb in token["rwkv_tokens"]]
-            model_a_info = base64.b64encode(json.dumps(model_a_list, ensure_ascii=False).encode("utf-8")).decode("ascii")
-
-        # Model B (Qwen3) - tokens overlapping this byte range
-        model_b_info = ""
-        if token["qwen_tokens"]:
-            model_b_list = [[tid, *token_bytes_to_display_text(tb)] for tid, tb in token["qwen_tokens"]]
-            model_b_info = base64.b64encode(json.dumps(model_b_list, ensure_ascii=False).encode("utf-8")).decode("ascii")
-
         raw_bytes = list(text_bytes[byte_start:byte_end])
         losses_a = byte_losses_a[byte_start:byte_end]
         losses_b = byte_losses_b[byte_start:byte_end]
@@ -426,8 +490,8 @@ def generate_comparison_html(
         avg_compression_a_token = sum(losses_a) / len(losses_a) * COMPRESSION_RATE_FACTOR if losses_a else 0
         avg_compression_b_token = sum(losses_b) / len(losses_b) * COMPRESSION_RATE_FACTOR if losses_b else 0
 
-        topk_a_json = ""
-        topk_b_json = ""
+        topk_a_data = None
+        topk_b_data = None
         if topk_predictions_a is not None and model_a_token_ranges:
             model_a_token_idx = find_token_for_byte(byte_start, model_a_token_ranges)
             if model_a_token_idx is not None and model_a_token_idx < len(topk_predictions_a):
@@ -435,19 +499,18 @@ def generate_comparison_html(
                 try:
                     if len(pred) >= 4:
                         actual_id, rank, actual_prob, topk_list = pred[0], pred[1], pred[2], pred[3]
-                        decoded_pred = [
+                        topk_a_data = [
                             actual_id,
                             rank,
                             actual_prob,
                             [[tid, prob, *decode_token(tid, tokenizer_a, model_type_a)] for tid, prob in topk_list],
                         ]
                     else:
-                        decoded_pred = [
+                        topk_a_data = [
                             pred[0],
                             pred[1],
                             [[tid, prob, *decode_token(tid, tokenizer_a, model_type_a)] for tid, prob in pred[2]],
                         ]
-                    topk_a_json = base64.b64encode(json.dumps(decoded_pred, ensure_ascii=False).encode("utf-8")).decode("ascii")
                 except Exception as e:
                     pass
         if topk_predictions_b is not None and model_b_token_ranges:
@@ -457,20 +520,17 @@ def generate_comparison_html(
                 try:
                     if len(pred) >= 4:
                         actual_id, rank, actual_prob, topk_list = pred[0], pred[1], pred[2], pred[3]
-                        decoded_pred = [
+                        topk_b_data = [
                             actual_id,
                             rank,
                             actual_prob,
                             [[tid, prob, *decode_token(tid, tokenizer_b, model_type_b)] for tid, prob in topk_list],
                         ]
                     else:
-                        decoded_pred = [pred[0], pred[1], [[tid, prob, *decode_token(tid, tokenizer_b, model_type_b)] for tid, prob in pred[2]]]
-                    topk_b_json = base64.b64encode(json.dumps(decoded_pred, ensure_ascii=False).encode("utf-8")).decode("ascii")
+                        topk_b_data = [pred[0], pred[1], [[tid, prob, *decode_token(tid, tokenizer_b, model_type_b)] for tid, prob in pred[2]]]
                 except Exception as e:
                     pass
 
-        token_count += 1
-
         token_deltas = deltas[byte_start:byte_end]
         avg_token_delta = sum(token_deltas) / len(token_deltas) if token_deltas else 0
         tuned_delta = avg_token_delta - avg_delta
@@ -478,695 +538,124 @@ def generate_comparison_html(
         # Initial rendering uses white color, JavaScript will apply colors based on slider
         r, g, b = 255, 255, 255
 
-        token_html_parts = []
-        for char in token_text:
-            if char == "<":
-                escaped_char = "&lt;"
-            elif char == ">":
-                escaped_char = "&gt;"
-            elif char == "&":
-                escaped_char = "&amp;"
-            elif char == "\t":
-                escaped_char = "    "
-            else:
-                escaped_char = char
-            token_html_parts.append(escaped_char)
-
-        token_span_content = "".join(token_html_parts)
-        data_attrs = (
-            f'data-model-a="{escape_for_attr(model_a_info)}" '
-            f'data-model-b="{escape_for_attr(model_b_info)}" '
-            f'data-bytes="{escape_for_attr(bytes_str)}" '
-            f'data-compression-a="{escape_for_attr(compression_a_str)}" '
-            f'data-compression-b="{escape_for_attr(compression_b_str)}" '
-            f'data-avg-compression-a="{avg_compression_a_token:.2f}" '
-            f'data-avg-compression-b="{avg_compression_b_token:.2f}" '
-            f'data-tuned-delta="{tuned_delta:.6f}" '
-            f'data-topk-a="{escape_for_attr(topk_a_json)}" '
-            f'data-topk-b="{escape_for_attr(topk_b_json)}"'
-        )
-        style_attr = f'style="background-color: rgb({r},{g},{b})"'
+        raw_display_text = token_text
+        display_text = token_text.replace("\t", "    ")
 
-        if token["type"] == "word":
-            word_lower = token["word_lower"]
-            occurrences = word_occurrences[word_lower]
-            if len(occurrences) > 1:
-                word_id = token["word_id"]
-                html_content.append(
-                    f'<span class="token word" {data_attrs} {style_attr} data-word="{word_lower}" data-word-id="{word_id}">'
-                    + token_span_content
-                    + "</span>"
-                )
-            else:
-                html_content.append(f'<span class="token" {data_attrs} {style_attr}>{token_span_content}</span>')
-        else:
-            html_content.append(f'<span class="token" {data_attrs} {style_attr}>{token_span_content}</span>')
+        def classify_kind(text_value: str, is_raw_value: bool) -> str:
+            return build_display(text_value, is_raw=is_raw_value).kind
+
+        def get_actual_prob(topk_predictions, token_idx: Optional[int]):
+            if not topk_predictions or token_idx is None:
+                return None
+            if token_idx < 0 or token_idx >= len(topk_predictions):
+                return None
+            pred = topk_predictions[token_idx]
+            if isinstance(pred, (list, tuple)) and len(pred) >= 3:
+                return pred[2]
+            return None
+
+        model_tokens_render = {}
+        if token["rwkv_tokens"]:
+            rwkv_items = []
+            for tok_idx, tid, tb in token["rwkv_tokens"]:
+                txt, is_raw = token_bytes_to_display_text(tb)
+                rwkv_items.append([tid, txt, classify_kind(txt, is_raw), get_actual_prob(topk_predictions_a, tok_idx)])
+            model_tokens_render["rwkv"] = rwkv_items
+        if token["qwen_tokens"]:
+            qwen_items = []
+            for tok_idx, tid, tb in token["qwen_tokens"]:
+                txt, is_raw = token_bytes_to_display_text(tb)
+                qwen_items.append([tid, txt, classify_kind(txt, is_raw), get_actual_prob(topk_predictions_b, tok_idx)])
+            model_tokens_render["qwen"] = qwen_items
+
+        display_info = build_display(raw_display_text, is_raw=not decoded_ok)
+        if display_info.kind == "control":
+            display_text = raw_display_text
+        display_info.text = display_text
+        render_tokens.append(
+            TokenInfo(
+                byte_start=byte_start,
+                byte_end=byte_end,
+                display=display_info,
+                is_word=token["type"] == "word",
+                word_id=token.get("word_id"),
+                word_key=token.get("word_lower"),
+                bytes_hex=bytes_str,
+                compression={"rwkv": compression_a_str, "qwen": compression_b_str},
+                model_tokens=model_tokens_render,
+                loss={"rwkv": avg_compression_a_token, "qwen": avg_compression_b_token},
+                topk={
+                    "rwkv": topk_a_data,
+                    "qwen": topk_b_data,
+                },
+                tuned_delta=tuned_delta,
+            )
+        )
 
     delta_color = "#64ff64" if avg_delta < 0 else "#ff6464"
 
-    html = f"""<!DOCTYPE html>
-<html>
-<head>
-    <meta charset="UTF-8">
-    <title>Model Comparison</title>
-    <style>
-        body {{
-            font-family: Consolas, 'Courier New', monospace;
-            margin: 0;
-            padding: 0;
-            background-color: #f5f5f5;
-        }}
-        .header {{
-            background-color: #333;
-            color: white;
-            padding: 20px;
-            position: sticky;
-            top: 0;
-            z-index: 100;
-        }}
-        .header h1 {{
-            margin: 0 0 15px 0;
-            font-size: 18px;
-        }}
-        .meta {{
-            display: flex;
-            flex-wrap: wrap;
-            gap: 20px;
-            font-size: 12px;
-            color: #c8c8c8;
-        }}
-        .legend {{
-            display: flex;
-            gap: 15px;
-            margin-top: 10px;
-        }}
-        .legend-item {{
-            display: flex;
-            align-items: center;
-            gap: 5px;
-        }}
-        .legend-box {{
-            width: 20px;
-            height: 12px;
-            border: 1px solid #666;
-        }}
-        .content {{
-            background-color: white;
-            margin: 10px;
-            padding: 15px;
-            border: 1px solid #ccc;
-            font-size: 14px;
-            line-height: 1.8;
-            white-space: pre-wrap;
-            overflow-wrap: anywhere;
-            word-wrap: break-word;
-            position: relative;
-        }}
-        .content span {{
-            padding: 1px 0;
-        }}
-        .word {{
-            cursor: pointer;
-            position: relative;
-        }}
-        .word:hover {{
-            outline: 2px solid #007bff;
-            outline-offset: 1px;
-        }}
-        .word.highlighted {{
-            outline: 2px solid #ff6b6b;
-            outline-offset: 1px;
-        }}
-        #svg-overlay {{
-            position: fixed;
-            top: 0;
-            left: 0;
-            width: 100%;
-            height: 100%;
-            pointer-events: none;
-            z-index: 1000;
-        }}
-        .link-line {{
-            stroke: #007bff;
-            stroke-width: 2;
-            fill: none;
-            opacity: 0.7;
-        }}
-        .link-dot {{
-            fill: #007bff;
-            opacity: 0.8;
-        }}
-        .token {{
-            position: relative;
-            cursor: help;
-        }}
-        .token:hover {{
-            outline: 1px dashed #666;
-        }}
-        #tooltip {{
-            position: fixed;
-            background-color: rgba(0, 0, 0, 0.9);
-            color: white;
-            padding: 10px 14px;
-            border-radius: 6px;
-            font-size: 12px;
-            max-width: 500px;
-            z-index: 2000;
-            pointer-events: none;
-            display: none;
-            line-height: 1.6;
-            box-shadow: 0 2px 10px rgba(0,0,0,0.3);
-        }}
-        #tooltip .label {{
-            color: #aaa;
-            font-weight: bold;
-        }}
-        #tooltip .bytes {{
-            color: #a5f3fc;
-            font-family: monospace;
-        }}
-        #tooltip .loss-a {{
-            color: #86efac;
-            font-family: monospace;
-        }}
-        #tooltip .loss-b {{
-            color: #fca5a5;
-            font-family: monospace;
-        }}
-        #tooltip .model-a {{
-            color: #fcd34d;
-        }}
-        #tooltip .model-b {{
-            color: #7dd3fc;
-        }}
-        #tooltip .topk-section {{
-            margin-top: 8px;
-            padding-top: 8px;
-            border-top: 1px solid #555;
-        }}
-        #tooltip .topk-container {{
-            display: flex;
-            gap: 16px;
-        }}
-        #tooltip .topk-column {{
-            flex: 1;
-            min-width: 180px;
-        }}
-        #tooltip .topk-title {{
-            color: #aaa;
-            font-weight: bold;
-            margin-bottom: 4px;
-            font-size: 11px;
-        }}
-        #tooltip .topk-title.model-a {{
-            color: #86efac;
-        }}
-        #tooltip .topk-title.model-b {{
-            color: #fca5a5;
-        }}
-        #tooltip .topk-list {{
-            font-size: 11px;
-        }}
-        #tooltip .topk-item {{
-            display: flex;
-            gap: 4px;
-            padding: 1px 0;
-            align-items: flex-start;
-        }}
-        #tooltip .token-block {{
-            margin-top: 6px;
-            display: flex;
-            align-items: center;
-            gap: 6px;
-            white-space: nowrap;
-        }}
-        #tooltip .token-chips {{
-            display: flex;
-            flex-wrap: nowrap;
-            gap: 4px;
-        }}
-        #tooltip .token-chip-group {{
-            display: inline-flex;
-            align-items: center;
-            gap: 4px;
-        }}
-        #tooltip .token-id {{
-            color: #888;
-            font-family: monospace;
-        }}
-        #tooltip .token-chip {{
-            max-width: 100%;
-        }}
-        #tooltip .topk-rank {{
-            color: #888;
-            min-width: 18px;
-        }}
-        #tooltip .topk-rank.hit {{
-            color: #ffd700;
-        }}
-        #tooltip .topk-token {{
-            color: #a5f3fc;
-            white-space: pre-wrap;
-            overflow-wrap: anywhere;
-            word-break: break-word;
-            font-family: monospace;
-            background-color: rgba(255, 255, 255, 0.08);
-            padding: 0 4px;
-            border-radius: 3px;
-            display: inline-block;
-            max-width: 100%;
-        }}
-        #tooltip .esc-control {{
-            color: #fbbf24;
-        }}
-        #tooltip .esc-raw {{
-            color: #fb7185;
-        }}
-        #tooltip .topk-prob {{
-            color: #86efac;
-            min-width: 45px;
-            text-align: right;
-        }}
-        #tooltip .topk-hit {{
-            color: #22c55e;
-        }}
-        #tooltip .topk-miss {{
-            color: #ef4444;
-            font-style: italic;
-        }}
-    </style>
-</head>
-<body>
-    <svg id="svg-overlay"></svg>
-    <div id="tooltip"></div>
-    <div class="header">
-        <div class="meta">
-            <div>Model A: {model_a_name}</div>
-            <div>Model B: {model_b_name}</div>
-            <div>RWKV Compression: {avg_compression_a:.2f}%</div>
-            <div>Qwen Compression: {avg_compression_b:.2f}%</div>
-            <div style="color: {delta_color}">Avg Delta: {avg_delta_compression:+.2f}%</div>
-        </div>
-        <div class="legend">
-            <div class="legend-item">
-                <div class="legend-box" style="background-color: rgb(77, 255, 77)"></div>
-                <span>RWKV better than avg</span>
-            </div>
-            <div class="legend-item">
-                <div class="legend-box" style="background-color: rgb(255, 255, 255)"></div>
-                <span>Equal to avg</span>
+    render_model = RenderModel(
+        text=text,
+        tokens=render_tokens,
+        meta={
+            "model_a": model_a_name,
+            "model_b": model_b_name,
+            "avg_compression": {
+                "rwkv": avg_compression_a,
+                "qwen": avg_compression_b,
+            },
+            "avg_delta": avg_delta,
+            "avg_delta_compression": avg_delta_compression,
+        },
+    )
+    render_model_json = escape_json_for_script(render_model.to_dict())
+
+    style_block = (ASSETS_DIR / "main.css").read_text(encoding="utf-8")
+
+    header_html = f"""
+        <div class="header">
+            <div class="meta">
+                <div>Model A: {model_a_name}</div>
+                <div>Model B: {model_b_name}</div>
+                <div>RWKV Compression: {avg_compression_a:.2f}%</div>
+                <div>Qwen Compression: {avg_compression_b:.2f}%</div>
+                <div style="color: {delta_color}">Avg Delta: {avg_delta_compression:+.2f}%</div>
             </div>
-            <div class="legend-item">
-                <div class="legend-box" style="background-color: rgb(255, 77, 77)"></div>
-                <span>RWKV worse than avg</span>
-            </div>
-            <div class="legend-item" style="margin-left: 20px;">
-                <span style="color: #aaa;">Color Range:</span>
-                <input type="range" id="color-range-slider" min="0" max="100" value="10" step="0.1" style="width: 200px; vertical-align: middle;">
-                <span id="color-range-value" style="color: #fff; min-width: 45px; display: inline-block;">10%</span>
+            <div class="legend">
+                <div class="legend-item">
+                    <div class="legend-box" style="background-color: rgb(77, 255, 77)"></div>
+                    <span>RWKV better than avg</span>
+                </div>
+                <div class="legend-item">
+                    <div class="legend-box" style="background-color: rgb(255, 255, 255)"></div>
+                    <span>Equal to avg</span>
+                </div>
+                <div class="legend-item">
+                    <div class="legend-box" style="background-color: rgb(255, 77, 77)"></div>
+                    <span>RWKV worse than avg</span>
+                </div>
+                <div class="legend-item" style="margin-left: 20px;">
+                    <span style="color: #aaa;">Color Range:</span>
+                    <input type="range" id="color-range-slider" min="0" max="100" value="10" step="0.1" style="width: 200px; vertical-align: middle;">
+                    <span id="color-range-value" style="color: #fff; min-width: 45px; display: inline-block;">10%</span>
+                </div>
             </div>
         </div>
-    </div>
-    <div class="content">{''.join(html_content)}</div>
-    <script>
-        const svgOverlay = document.getElementById('svg-overlay');
-        const words = document.querySelectorAll('.word');
-
-        const wordGroups = {{}};
-        words.forEach(word => {{
-            const wordText = word.getAttribute('data-word');
-            if (!wordGroups[wordText]) {{
-                wordGroups[wordText] = [];
-            }}
-            wordGroups[wordText].push(word);
-        }});
-
-        function clearLines() {{
-            svgOverlay.innerHTML = '';
-            words.forEach(w => w.classList.remove('highlighted'));
-        }}
-
-        function pickRectByY(rects, targetY) {{
-            if (!rects || rects.length === 0) return null;
-            let best = rects[0];
-            let bestDist = Infinity;
-            rects.forEach(r => {{
-                const cy = r.top + r.height / 2;
-                const dist = Math.abs(cy - targetY);
-                if (dist < bestDist) {{
-                    best = r;
-                    bestDist = dist;
-                }}
-            }});
-            return best;
-        }}
-
-        function getAnchorRect(element, targetY) {{
-            const rects = Array.from(element.getClientRects());
-            if (rects.length === 0) return element.getBoundingClientRect();
-            if (rects.length === 1) return rects[0];
-            const picked = pickRectByY(rects, targetY);
-            return picked || rects[0];
-        }}
-
-        function drawLines(hoveredWord, evt) {{
-            clearLines();
-
-            const wordText = hoveredWord.getAttribute('data-word');
-            const wordId = parseInt(hoveredWord.getAttribute('data-word-id'));
-            const sameWords = wordGroups[wordText] || [];
-
-            const previousWords = sameWords.filter(w => {{
-                const id = parseInt(w.getAttribute('data-word-id'));
-                return id < wordId;
-            }});
-
-            if (previousWords.length === 0) return;
-
-            sameWords.forEach(w => w.classList.add('highlighted'));
-
-            const targetY = evt ? evt.clientY : (hoveredWord.getBoundingClientRect().top + hoveredWord.getBoundingClientRect().height / 2);
-            const hoveredRect = getAnchorRect(hoveredWord, targetY);
-            const hoveredX = hoveredRect.left + hoveredRect.width / 2;
-            const hoveredY = hoveredRect.top + hoveredRect.height / 2;
-
-            previousWords.forEach(prevWord => {{
-                const prevRect = getAnchorRect(prevWord, hoveredY);
-                const prevX = prevRect.left + prevRect.width / 2;
-                const prevY = prevRect.top + prevRect.height / 2;
-
-                const midX = (hoveredX + prevX) / 2;
-                const midY = Math.min(hoveredY, prevY) - 30;
-
-                const path = document.createElementNS('http://www.w3.org/2000/svg', 'path');
-                path.setAttribute('class', 'link-line');
-                path.setAttribute('d', `M ${{prevX}} ${{prevY}} Q ${{midX}} ${{midY}} ${{hoveredX}} ${{hoveredY}}`);
-                svgOverlay.appendChild(path);
-
-                const dot1 = document.createElementNS('http://www.w3.org/2000/svg', 'circle');
-                dot1.setAttribute('class', 'link-dot');
-                dot1.setAttribute('cx', prevX);
-                dot1.setAttribute('cy', prevY);
-                dot1.setAttribute('r', 4);
-                svgOverlay.appendChild(dot1);
-
-                const dot2 = document.createElementNS('http://www.w3.org/2000/svg', 'circle');
-                dot2.setAttribute('class', 'link-dot');
-                dot2.setAttribute('cx', hoveredX);
-                dot2.setAttribute('cy', hoveredY);
-                dot2.setAttribute('r', 4);
-                svgOverlay.appendChild(dot2);
-            }});
-        }}
-
-        words.forEach(word => {{
-            word.addEventListener('mouseenter', (e) => drawLines(word, e));
-            word.addEventListener('mouseleave', clearLines);
-        }});
-
-        window.addEventListener('scroll', clearLines);
-
-        const tooltip = document.getElementById('tooltip');
-        const tokenSpans = document.querySelectorAll('.token');
-
-        tokenSpans.forEach(token => {{
-            token.addEventListener('mouseenter', (e) => {{
-                const modelA = token.getAttribute('data-model-a') || '';
-                const modelB = token.getAttribute('data-model-b') || '';
-                const bytes = token.getAttribute('data-bytes') || '';
-                const compressionA = token.getAttribute('data-compression-a') || '';
-                const compressionB = token.getAttribute('data-compression-b') || '';
-                const avgCompressionA = token.getAttribute('data-avg-compression-a') || '';
-                const avgCompressionB = token.getAttribute('data-avg-compression-b') || '';
-                const top5A = token.getAttribute('data-topk-a') || '';
-                const top5B = token.getAttribute('data-topk-b') || '';
-
-                function decodeBase64Json(base64Str) {{
-                    const binaryString = atob(base64Str);
-                    const bytes = new Uint8Array(binaryString.length);
-                    for (let i = 0; i < binaryString.length; i++) {{
-                        bytes[i] = binaryString.charCodeAt(i);
-                    }}
-                    const jsonStr = new TextDecoder('utf-8').decode(bytes);
-                    return JSON.parse(jsonStr);
-                }}
-
-                function escapeControlChars(text) {{
-                    if (!text) return text;
-                    let out = '';
-                    for (let i = 0; i < text.length; i++) {{
-                        const ch = text[i];
-                        const code = text.charCodeAt(i);
-                        if (ch === '\\\\') {{
-                            out += '\\\\\\\\';
-                        }} else if (ch === '\\n') {{
-                            out += '\\\\n';
-                        }} else if (ch === '\\r') {{
-                            out += '\\\\r';
-                        }} else if (ch === '\\t') {{
-                            out += '\\\\t';
-                        }} else if (code < 32 || code === 127) {{
-                            out += '\\\\x' + code.toString(16).padStart(2, '0');
-                        }} else {{
-                            out += ch;
-                        }}
-                    }}
-                    return out;
-                }}
-
-                function renderEscapedWithControlColor(text) {{
-                    const escaped = (text || '')
-                        .replace(/&/g, '&amp;')
-                        .replace(/</g, '&lt;')
-                        .replace(/>/g, '&gt;');
-                    return escaped.replace(/\\\\(x[0-9a-fA-F]{2}|[nrt])/g, '<span class="esc-control">\\\\$1</span>');
-                }}
-
-                function formatTopkColumn(topkBase64, modelName, titleClass) {{
-                    if (!topkBase64) return '<div class="topk-column"><div class="topk-title ' + titleClass + '">' + modelName + '</div><div class="topk-list">N/A</div></div>';
-                    try {{
-                        const data = decodeBase64Json(topkBase64);
-                        let actualId = null;
-                        let rank = null;
-                        let actualProb = null;
-                        let topkList = [];
-                        if (data.length >= 4) {{
-                            [actualId, rank, actualProb, topkList] = data;
-                        }} else {{
-                            [actualId, rank, topkList] = data;
-                        }}
-                        let html = '<div class="topk-column">';
-                        html += '<div class="topk-title ' + titleClass + '">' + modelName + '</div>';
-                        html += '<div class="topk-list">';
-                        topkList.forEach((item, idx) => {{
-                            const tokenId = item[0];
-                            const prob = item[1];
-                            const tokenText = item[2];
-                            const isRaw = item.length > 3 ? item[3] : false;
-                            const isHit = tokenId === actualId;
-                            const rankClass = isHit ? 'topk-rank hit' : 'topk-rank';
-                            const rawText = (tokenText !== undefined && tokenText !== null) ? tokenText : '';
-                            let displayText = '';
-                            let htmlText = '';
-                            if (isRaw) {{
-                                displayText = (rawText !== '') ? rawText : ('[' + tokenId + ']');
-                                const escapedText = displayText
-                                    .replace(/&/g, '&amp;')
-                                    .replace(/</g, '&lt;')
-                                    .replace(/>/g, '&gt;');
-                                htmlText = '<span class="esc-raw">' + escapedText + '</span>';
-                            }} else {{
-                                const visibleText = escapeControlChars(rawText);
-                                displayText = (visibleText !== '') ? visibleText : ('[' + tokenId + ']');
-                                htmlText = renderEscapedWithControlColor(displayText);
-                            }}
-                            html += '<div class="topk-item">';
-                            html += '<span class="' + rankClass + '">' + (idx + 1) + '.</span>';
-                            html += '<span class="topk-token" title="ID: ' + tokenId + '">' + htmlText + '</span>';
-                            html += '<span class="topk-prob">' + (prob * 100).toFixed(1) + '%</span>';
-                            if (isHit) html += '<span class="topk-hit">✓</span>';
-                            html += '</div>';
-                        }});
-                        if (rank > 10) {{
-                            let probSuffix = '';
-                            const probVal = parseFloat(actualProb);
-                            if (!isNaN(probVal)) {{
-                                probSuffix = ' (' + (probVal * 100).toFixed(4) + '%)';
-                            }}
-                            html += '<div class="topk-item topk-miss">Actual rank: ' + rank + probSuffix + '</div>';
-                        }}
-                        html += '</div></div>';
-                        return html;
-                    }} catch (e) {{
-                        console.error('Error in formatTopkColumn for ' + modelName + ':', e);
-                        console.error('topkBase64:', topkBase64);
-                        return '<div class="topk-column"><div class="topk-title ' + titleClass + '">' + modelName + '</div><div class="topk-list">Error: ' + e.message + '</div></div>';
-                    }}
-                }}
-
-                function formatTokenChips(modelBase64, label, labelClass) {{
-                    if (!modelBase64) {{
-                        return '<div class="token-block"><span class="label ' + labelClass + '">' + label + ':</span> <span class="topk-token token-chip">N/A</span></div>';
-                    }}
-                    try {{
-                        const tokenList = decodeBase64Json(modelBase64);
-                        let html = '<div class="token-block">';
-                        html += '<span class="label ' + labelClass + '">' + label + ':</span>';
-                        html += '<div class="token-chips">';
-                        tokenList.forEach((item) => {{
-                            const tokenId = item[0];
-                            const tokenText = item[1];
-                            const isRaw = item.length > 2 ? item[2] : false;
-                            let displayText = '';
-                            let htmlText = '';
-                            if (isRaw) {{
-                                displayText = tokenText || '';
-                                const escapedText = displayText
-                                    .replace(/&/g, '&amp;')
-                                    .replace(/</g, '&lt;')
-                                    .replace(/>/g, '&gt;');
-                                htmlText = '<span class="esc-raw">' + escapedText + '</span>';
-                            }} else {{
-                                const visible = escapeControlChars(tokenText || '');
-                                displayText = (visible !== '') ? visible : '';
-                                htmlText = renderEscapedWithControlColor(displayText);
-                            }}
-                            html += '<span class="token-chip-group" title="ID: ' + tokenId + '">';
-                            html += '<span class="token-id">[' + tokenId + ']</span>';
-                            html += '<span class="topk-token token-chip">' + htmlText + '</span>';
-                            html += '</span>';
-                        }});
-                        html += '</div></div>';
-                        return html;
-                    }} catch (e) {{
-                        console.error('Error in formatTokenChips for ' + label + ':', e);
-                        console.error('modelBase64:', modelBase64);
-                        return '<div class="token-block"><span class="label ' + labelClass + '">' + label + ':</span> <span class="topk-token token-chip">Error: ' + e.message + '</span></div>';
-                    }}
-                }}
-
-                let tooltipHtml = `
-                    <div><span class="label">Bytes:</span> <span class="bytes">${{bytes || '(empty)'}}</span></div>
-                    <div><span class="label">RWKV Compression Rate:</span> <span class="loss-a">${{compressionA || '(empty)'}}${{avgCompressionA ? ' (avg: ' + avgCompressionA + '%)' : ''}}</span></div>
-                    <div><span class="label">Qwen Compression Rate:</span> <span class="loss-b">${{compressionB || '(empty)'}}${{avgCompressionB ? ' (avg: ' + avgCompressionB + '%)' : ''}}</span></div>
-                    <hr style="border-color: #555; margin: 6px 0;">
-                    ${{formatTokenChips(modelA, 'RWKV', 'model-a')}}
-                    ${{formatTokenChips(modelB, 'Qwen', 'model-b')}}
-                `;
-                if (top5A || top5B) {{
-                    tooltipHtml += '<div class="topk-section"><div class="topk-container">';
-                    tooltipHtml += formatTopkColumn(top5A, 'RWKV Top10', 'model-a');
-                    tooltipHtml += formatTopkColumn(top5B, 'Qwen Top10', 'model-b');
-                    tooltipHtml += '</div></div>';
-                }}
-                tooltip.innerHTML = tooltipHtml;
-                tooltip.style.display = 'block';
-            }});
-
-            token.addEventListener('mousemove', (e) => {{
-                const tooltipRect = tooltip.getBoundingClientRect();
-                const viewportWidth = window.innerWidth;
-                const viewportHeight = window.innerHeight;
-
-                let x = e.clientX + 15;
-                let y = e.clientY + 15;
-
-                if (x + tooltipRect.width > viewportWidth - 10) {{
-                    x = e.clientX - tooltipRect.width - 15;
-                }}
-                if (y + tooltipRect.height > viewportHeight - 10) {{
-                    y = e.clientY - tooltipRect.height - 15;
-                }}
-                if (x < 10) x = 10;
-                if (y < 10) y = 10;
-
-                tooltip.style.left = x + 'px';
-                tooltip.style.top = y + 'px';
-            }});
-
-            token.addEventListener('mouseleave', () => {{
-                tooltip.style.display = 'none';
-            }});
-        }});
-
-        const slider = document.getElementById('color-range-slider');
-        const rangeValue = document.getElementById('color-range-value');
-
-        // Collect all tuned_delta values
-        const tokenData = [];
-        tokenSpans.forEach((token, idx) => {{
-            const tunedDelta = parseFloat(token.getAttribute('data-tuned-delta'));
-            if (!isNaN(tunedDelta)) {{
-                tokenData.push({{ token, tunedDelta, absDelta: Math.abs(tunedDelta) }});
-            }}
-        }});
-
-        // Calculate max_abs_tuned_delta for normalization
-        const maxAbsDelta = Math.max(...tokenData.map(d => d.absDelta), 1e-9);
-
-        // Sort by |tuned_delta| to get rankings
-        const sortedByAbs = [...tokenData].sort((a, b) => b.absDelta - a.absDelta);
-        sortedByAbs.forEach((item, rank) => {{
-            item.rank = rank;  // rank 0 = largest deviation
-        }});
-
-        function tunedDeltaToColor(tunedDelta, maxAbsDelta, exponent) {{
-            // Normalize to [-1, 1]
-            const normalized = Math.max(-1, Math.min(1, tunedDelta / maxAbsDelta));
-            let r, g, b;
-            if (normalized < 0) {{
-                // Green (RWKV better)
-                const intensity = Math.pow(-normalized, exponent);
-                r = Math.round(255 * (1 - intensity * 0.85));
-                g = 255;
-                b = Math.round(255 * (1 - intensity * 0.85));
-            }} else {{
-                // Red (RWKV worse)
-                const intensity = Math.pow(normalized, exponent);
-                r = 255;
-                g = Math.round(255 * (1 - intensity * 0.85));
-                b = Math.round(255 * (1 - intensity * 0.85));
-            }}
-            return `rgb(${{r}}, ${{g}}, ${{b}})`;
-        }}
-
-        function updateColors(colorRangePercent) {{
-            // colorRangePercent: 0-100, represents the proportion of tokens to color
-            const colorCount = Math.round(tokenData.length * colorRangePercent / 100);
-
-            // Calculate exponent: 100% -> 0.5, 0% -> 1.0
-            const exponent = 1 - (colorRangePercent / 100) * 0.5;
-
-            // Calculate max deviation within the colored range
-            let maxAbsDeltaInRange = 1e-9;
-            tokenData.forEach(item => {{
-                if (item.rank < colorCount) {{
-                    maxAbsDeltaInRange = Math.max(maxAbsDeltaInRange, item.absDelta);
-                }}
-            }});
-
-            tokenData.forEach(item => {{
-                if (item.rank < colorCount) {{
-                    // Use dynamic normalization based on colored range
-                    item.token.style.backgroundColor = tunedDeltaToColor(item.tunedDelta, maxAbsDeltaInRange, exponent);
-                }} else {{
-                    // Outside color range, white
-                    item.token.style.backgroundColor = 'rgb(255, 255, 255)';
-                }}
-            }});
-        }}
-
-        slider.addEventListener('input', (e) => {{
-            const val = parseFloat(e.target.value);
-            rangeValue.textContent = val.toFixed(1) + '%';
-            updateColors(val);
-        }});
-
-        // Apply default color range on page load
-        updateColors(10);
-    </script>
-</body>
-</html>
-"""
-
-    return html
+    """.strip("\n")
+
+    script_body = (ASSETS_DIR / "main.js").read_text(encoding="utf-8")
+
+    html_doc = render_page(
+        {
+            "page_title": "Model Comparison",
+            "style_block": style_block.strip("\n"),
+            "header_html": header_html,
+            "content_html": "",
+            "render_model_json": render_model_json,
+            "script_body": script_body.strip("\n"),
+        }
+    )
+
+    if return_render_model:
+        return html_doc, render_model.to_dict()
+    return html_doc

visualization/render.py

@@ -0,0 +1,16 @@
+"""
+Page rendering using a lightweight template.
+"""
+
+from pathlib import Path
+from string import Template
+from typing import Dict
+
+
+_TEMPLATE_PATH = Path(__file__).resolve().parent / "templates" / "page.html.tmpl"
+
+
+def render_page(context: Dict[str, str]) -> str:
+    template_text = _TEMPLATE_PATH.read_text(encoding="utf-8")
+    template = Template(template_text)
+    return template.safe_substitute(context)

visualization/templates/page.html.tmpl

@@ -0,0 +1,20 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <meta charset="UTF-8">
+    <title>$page_title</title>
+    <style>
+$style_block
+    </style>
+</head>
+<body>
+    <svg id="svg-overlay"></svg>
+    <div id="tooltip"></div>
+$header_html
+    <div class="content">$content_html</div>
+    <script id="render-model" type="application/json">$render_model_json</script>
+    <script>
+$script_body
+    </script>
+</body>
+</html>