Llama.cpp Tokenize Visualizer

tokenizer_viz.py

#!/usr/bin/env python3
"""
llama.cpp トークナイザー可視化ツール
単一Pythonファイル版 - llama-tokenizeコマンドの結果を色付きで表示
"""

import os
import glob
import subprocess
import json
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import parse_qs, urlparse
import html

# ============================================================
# ★ 設定 - ここを環境に合わせて変更してください
# ============================================================
LLAMA_CPP_PATH = "PATH_TO_LLAMA_CPP"          # llama-tokenize があるディレクトリ
GGUF_MODELS_DIR = "PATH_TO_MODEL_DIR"         # GGUFモデルが入っているディレクトリ
PORT = 8080
# ============================================================

# Windows なら .exe を付ける
_bin_name = "llama-tokenize.exe" if os.name == "nt" else "llama-tokenize"
LLAMA_TOKENIZE_BIN = os.path.join(LLAMA_CPP_PATH, _bin_name)

# トークン色パレット（明るく視認しやすい20色）
TOKEN_COLORS = [
    "#FF6B6B", "#FF9F43", "#FECA57", "#48DBFB", "#FF9FF3",
    "#54A0FF", "#5F27CD", "#00D2D3", "#1DD1A1", "#C8D6E5",
    "#FD9644", "#A29BFE", "#6C5CE7", "#00CEC9", "#E17055",
    "#FDCB6E", "#81ECEC", "#74B9FF", "#FAB1A0", "#55EFC4",
]

def get_gguf_models():
    """GGUFモデルファイルの一覧を取得（mmproj除外・ファイル名順ソート）"""
    patterns = [
        os.path.join(GGUF_MODELS_DIR, "*.gguf"),
        os.path.join(GGUF_MODELS_DIR, "**", "*.gguf"),
    ]
    found = []
    for pattern in patterns:
        found.extend(glob.glob(pattern, recursive=True))
    # 重複除去・mmproj 除外・ファイル名（大文字小文字無視）でソート
    models = sorted(
        {m for m in found if not os.path.basename(m).lower().startswith("mmproj")},
        key=lambda p: os.path.basename(p).lower()
    )
    return models

def run_llama_tokenize(model_path: str, prompt: str):
    """llama-tokenizeを実行してトークン情報を返す"""
    if not os.path.isfile(LLAMA_TOKENIZE_BIN):
        return None, f"llama-tokenize が見つかりません: {LLAMA_TOKENIZE_BIN}"
    if not os.path.isfile(model_path):
        return None, f"モデルが見つかりません: {model_path}"

    cmd = [LLAMA_TOKENIZE_BIN, "-m", model_path, "-p", prompt, "--show-count"]
    try:
        result = subprocess.run(
            cmd, capture_output=True, timeout=30  # text=True を削除
        )
        # バイト列をUTF-8でデコード（不正バイトは置換文字?に変換）
        stdout = result.stdout.decode("utf-8", errors="replace")
        stderr = result.stderr.decode("utf-8", errors="replace")
        output = stdout + stderr
        return output, None
    except subprocess.TimeoutExpired:
        return None, "タイムアウト (30秒)"
    except Exception as e:
        return None, str(e)

def parse_tokens(output: str):
    """
    llama-tokenize の実際の出力形式をパース:
        31171 -> 'お'
       151798 -> 'っぱ'
        16098 -> 'い'
       Total number of tokens: 3
    """
    import re
    tokens = []
    total = None

    # トークン行: 先頭に空白があってもよい数字、' -> '、シングルクォートで囲まれた文字列
    # piece はクォート内に何でも入りうる（空文字、特殊文字含む）
    token_re = re.compile(r'^\s*(\d+)\s*->\s*\'(.*)\'\s*$')

    # Total 行のパターン（複数バリエーション対応）
    total_re = re.compile(r'(?:total number of tokens|total tokens)[:\s]+(\d+)', re.IGNORECASE)

    for line in output.splitlines():
        m = token_re.match(line)
        if m:
            token_id = int(m.group(1))
            piece = m.group(2)
            tokens.append({"index": len(tokens), "id": token_id, "piece": piece})
            continue

        m2 = total_re.search(line)
        if m2:
            total = int(m2.group(1))

    return tokens, total

def build_html(models, selected_model="", prompt="", raw_output="", tokens=None, total=None, error=None):
    model_options = ""
    for m in models:
        sel = 'selected' if m == selected_model else ''
        name = os.path.basename(m)
        model_options += f'<option value="{html.escape(m)}" {sel}>{html.escape(name)}</option>\n'

    # トークン可視化HTML生成
    token_viz_html = ""
    token_table_html = ""
    if tokens:
        spans = []
        for i, tok in enumerate(tokens):
            color = TOKEN_COLORS[i % len(TOKEN_COLORS)]
            piece_display = tok["piece"].replace("▁", " ").replace("Ġ", " ")
            piece_escaped = html.escape(piece_display)
            title = f"id={tok['id']} piece={html.escape(tok['piece'])}"
            spans.append(
                f'<span class="token" style="background:{color}" title="{title}" data-idx="{i}">'
                f'{piece_escaped}</span>'
            )
        token_viz_html = "".join(spans)

        rows = []
        for i, tok in enumerate(tokens):
            color = TOKEN_COLORS[i % len(TOKEN_COLORS)]
            piece_escaped = html.escape(tok["piece"])
            rows.append(
                f'<tr><td><span class="dot" style="background:{color}"></span>{tok["index"]}</td>'
                f'<td>{tok["id"]}</td>'
                f'<td class="piece-cell">{piece_escaped}</td></tr>'
            )
        token_table_html = "\n".join(rows)

    total_str = str(total) if total is not None else (str(len(tokens)) if tokens else "-")
    error_html = f'<div class="error">{html.escape(error)}</div>' if error else ""
    raw_html = html.escape(raw_output) if raw_output else ""

    has_result = bool(tokens or error or raw_output)
    result_display = "block" if has_result else "none"

    # キャプチャ対象パネル内に表示するモデル名・プロンプト
    cap_model = html.escape(os.path.basename(selected_model)) if selected_model else ""
    cap_prompt = html.escape(prompt)

    return f"""<!DOCTYPE html>
<html lang="ja">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>llama.cpp Tokenizer Visualizer</title>
<script src="https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.min.js"></script>
<style>
  @import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;600&family=Space+Grotesk:wght@300;400;600;700&display=swap');

  :root {{
    --bg: #0d0d0d;
    --surface: #161616;
    --surface2: #1e1e1e;
    --border: #2a2a2a;
    --text: #e8e8e0;
    --text-dim: #888;
    --accent: #c8ff00;
    --accent2: #ff6b6b;
    --mono: 'JetBrains Mono', monospace;
    --sans: 'Space Grotesk', sans-serif;
  }}

  * {{ box-sizing: border-box; margin: 0; padding: 0; }}

  body {{
    background: var(--bg);
    color: var(--text);
    font-family: var(--sans);
    min-height: 100vh;
    padding: 40px 20px;
  }}

  .container {{
    max-width: 960px;
    margin: 0 auto;
  }}

  header {{
    margin-bottom: 32px;
    border-bottom: 1px solid var(--border);
    padding-bottom: 20px;
    display: flex;
    align-items: center;
    justify-content: space-between;
    gap: 16px;
    flex-wrap: wrap;
  }}

  .header-left {{ display: flex; flex-direction: column; gap: 4px; }}

  .logo {{
    display: flex;
    align-items: baseline;
    gap: 10px;
  }}

  .logo h1 {{
    font-size: 1.8rem;
    font-weight: 700;
    letter-spacing: -0.03em;
    color: var(--text);
  }}

  .logo .badge {{
    background: var(--accent);
    color: #000;
    font-family: var(--mono);
    font-size: 0.65rem;
    font-weight: 600;
    padding: 2px 8px;
    border-radius: 2px;
    letter-spacing: 0.08em;
    text-transform: uppercase;
  }}

  .subtitle {{
    color: var(--text-dim);
    font-size: 0.9rem;
    font-weight: 300;
  }}

  /* ── フォームカード ── */
  .card {{
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: 24px 28px;
    margin-bottom: 24px;
  }}

  .card-title {{
    font-size: 0.7rem;
    font-family: var(--mono);
    color: var(--text-dim);
    letter-spacing: 0.12em;
    text-transform: uppercase;
    margin-bottom: 16px;
  }}

  /* フォームを横並びにする */
  .form-row {{
    display: flex;
    gap: 16px;
    align-items: flex-end;
    flex-wrap: wrap;
  }}

  .form-col {{
    display: flex;
    flex-direction: column;
    gap: 6px;
  }}

  .form-col.model-col {{ flex: 0 0 300px; min-width: 200px; }}
  .form-col.prompt-col {{ flex: 1 1 200px; }}

  label {{
    font-size: 0.75rem;
    color: var(--text-dim);
    font-family: var(--mono);
    letter-spacing: 0.05em;
  }}

  select, textarea, input[type=text] {{
    background: var(--surface2);
    border: 1px solid var(--border);
    color: var(--text);
    border-radius: 4px;
    font-family: var(--mono);
    font-size: 0.85rem;
    padding: 10px 14px;
    outline: none;
    transition: border-color 0.2s;
    appearance: none;
    width: 100%;
  }}

  select:focus, textarea:focus {{ border-color: var(--accent); }}

  select {{
    cursor: pointer;
    background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='12' height='8' viewBox='0 0 12 8'%3E%3Cpath d='M1 1l5 5 5-5' stroke='%23888' stroke-width='1.5' fill='none' stroke-linecap='round'/%3E%3C/svg%3E");
    background-repeat: no-repeat;
    background-position: right 14px center;
    padding-right: 36px;
  }}

  textarea {{
    min-height: 46px;
    max-height: 120px;
    resize: vertical;
    line-height: 1.6;
  }}

  /* 実行ボタン */
  .btn-run {{
    background: var(--accent);
    color: #000;
    border: none;
    border-radius: 4px;
    padding: 11px 28px;
    font-family: var(--sans);
    font-size: 0.9rem;
    font-weight: 700;
    cursor: pointer;
    letter-spacing: 0.03em;
    white-space: nowrap;
    transition: opacity 0.15s, transform 0.1s;
    align-self: flex-end;
  }}
  .btn-run:hover {{ opacity: 0.85; transform: translateY(-1px); }}
  .btn-run:active {{ transform: translateY(0); }}

  /* キャプチャボタン */
  .btn-capture {{
    background: transparent;
    color: var(--text-dim);
    border: 1px solid var(--border);
    border-radius: 4px;
    padding: 8px 18px;
    font-family: var(--mono);
    font-size: 0.78rem;
    cursor: pointer;
    letter-spacing: 0.05em;
    transition: color 0.15s, border-color 0.15s;
    display: flex;
    align-items: center;
    gap: 6px;
  }}
  .btn-capture:hover {{ color: var(--accent); border-color: var(--accent); }}

  /* ── 結果エリア ── */
  .results {{ display: {result_display}; }}

  .error {{
    background: rgba(255,107,107,0.12);
    border: 1px solid rgba(255,107,107,0.4);
    color: var(--accent2);
    border-radius: 4px;
    padding: 14px 18px;
    font-family: var(--mono);
    font-size: 0.82rem;
    margin-bottom: 20px;
    white-space: pre-wrap;
  }}

  /* ── キャプチャ対象パネル ── */
  #capture-panel {{
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: 24px 28px;
    margin-bottom: 24px;
  }}

  .capture-meta {{
    margin-bottom: 16px;
    padding-bottom: 12px;
    border-bottom: 1px solid var(--border);
    display: flex;
    flex-wrap: wrap;
    gap: 16px;
    align-items: baseline;
  }}

  .meta-item {{
    display: flex;
    align-items: baseline;
    gap: 6px;
  }}

  .meta-label {{
    font-family: var(--mono);
    font-size: 0.65rem;
    color: var(--text-dim);
    letter-spacing: 0.1em;
    text-transform: uppercase;
  }}

  .meta-value {{
    font-family: var(--mono);
    font-size: 0.82rem;
    color: var(--text);
  }}

  /* 可視化ヘッダー（タイトル + トークン数 横並び） */
  .viz-header {{
    display: flex;
    align-items: center;
    justify-content: space-between;
    margin-bottom: 14px;
    gap: 12px;
  }}

  .viz-title {{
    font-size: 0.7rem;
    font-family: var(--mono);
    color: var(--text-dim);
    letter-spacing: 0.12em;
    text-transform: uppercase;
  }}

  .token-count-badge {{
    font-family: var(--mono);
    font-size: 0.78rem;
    background: var(--surface2);
    border: 1px solid var(--border);
    border-radius: 4px;
    padding: 3px 12px;
    color: var(--accent);
    white-space: nowrap;
  }}

  .token-count-badge span {{
    color: var(--text-dim);
    font-size: 0.68rem;
    margin-right: 4px;
  }}

  .viz-box {{
    background: var(--surface2);
    border: 1px solid var(--border);
    border-radius: 6px;
    padding: 20px;
    line-height: 2.4;
    word-break: break-all;
    font-family: var(--mono);
    font-size: 0.92rem;
    min-height: 60px;
  }}

  .token {{
    display: inline;
    border-radius: 3px;
    padding: 2px 1px;
    cursor: default;
    color: #000;
    font-weight: 700;
    transition: filter 0.15s;
    white-space: pre-wrap;
  }}
  .token:hover {{ filter: brightness(1.15) drop-shadow(0 0 5px currentColor); }}

  /* ── トークン一覧テーブル ── */
  .card .token-table {{
    width: 100%;
    border-collapse: collapse;
    font-family: var(--mono);
    font-size: 0.82rem;
  }}

  .token-table th {{
    text-align: left;
    padding: 8px 12px;
    border-bottom: 1px solid var(--border);
    color: var(--text-dim);
    font-size: 0.7rem;
    letter-spacing: 0.1em;
    text-transform: uppercase;
  }}

  .token-table td {{
    padding: 7px 12px;
    border-bottom: 1px solid rgba(42,42,42,0.6);
    vertical-align: middle;
  }}

  .token-table tr:last-child td {{ border-bottom: none; }}
  .token-table tr:hover td {{ background: rgba(255,255,255,0.03); }}

  .dot {{
    display: inline-block;
    width: 10px;
    height: 10px;
    border-radius: 50%;
    margin-right: 8px;
    vertical-align: middle;
  }}

  .piece-cell {{ color: var(--accent); }}

  /* ── Raw output ── */
  .raw-output {{
    background: var(--surface2);
    border: 1px solid var(--border);
    border-radius: 6px;
    padding: 16px;
    font-family: var(--mono);
    font-size: 0.78rem;
    white-space: pre-wrap;
    color: var(--text-dim);
    max-height: 240px;
    overflow-y: auto;
    line-height: 1.6;
  }}

  details summary {{
    cursor: pointer;
    font-family: var(--mono);
    font-size: 0.72rem;
    color: var(--text-dim);
    letter-spacing: 0.08em;
    text-transform: uppercase;
    margin-bottom: 10px;
    user-select: none;
  }}
  details summary:hover {{ color: var(--text); }}

  /* キャプチャ中は一時的に背景を固定色にする */
  .capturing #capture-panel {{
    background: #161616 !important;
  }}

  ::-webkit-scrollbar {{ width: 6px; height: 6px; }}
  ::-webkit-scrollbar-track {{ background: var(--surface); }}
  ::-webkit-scrollbar-thumb {{ background: var(--border); border-radius: 3px; }}
</style>
</head>
<body>
<div class="container">
  <header>
    <div class="header-left">
      <div class="logo">
        <h1>Tokenizer Viz</h1>
        <span class="badge">llama.cpp</span>
      </div>
      <p class="subtitle">llama-tokenize の出力をトークン単位で色分け可視化</p>
    </div>
  </header>

  <!-- ── 設定フォーム ── -->
  <div class="card">
    <form method="POST" action="/">
      <div class="form-row">
        <div class="form-col model-col">
          <label>GGUF モデル</label>
          <select name="model">
            <option value="">-- モデルを選択 --</option>
            {model_options}
          </select>
        </div>
        <div class="form-col prompt-col">
          <label>プロンプト</label>
          <textarea name="prompt" placeholder="トークナイズしたいテキストを入力...">{html.escape(prompt)}</textarea>
        </div>
        <button type="submit" class="btn-run">▶ 実行</button>
      </div>
    </form>
  </div>

  <!-- ── 結果 ── -->
  <div class="results">
    {error_html}

    {"" if not tokens else f'''
    <!-- キャプチャ対象パネル -->
    <div id="capture-panel">
      <!-- メタ情報（モデル・プロンプト） -->
      <div class="capture-meta">
        <div class="meta-item">
          <span class="meta-label">Model</span>
          <span class="meta-value">{cap_model}</span>
        </div>
        <div class="meta-item">
          <span class="meta-label">Prompt</span>
          <span class="meta-value">{cap_prompt}</span>
        </div>
      </div>

      <!-- 可視化ヘッダー：タイトル ＋ トークン数バッジ -->
      <div class="viz-header">
        <span class="viz-title">トークン可視化</span>
        <span class="token-count-badge"><span>TOTAL</span>{total_str} tokens</span>
      </div>

      <!-- カラーバー -->
      <div class="viz-box">{token_viz_html}</div>
    </div>

    <!-- キャプチャボタン（パネル外） -->
    <div style="display:flex; justify-content:flex-end; margin-top:-12px; margin-bottom:24px;">
      <button class="btn-capture" onclick="capturePanel()">
        <svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
          <rect x="3" y="3" width="18" height="18" rx="2"/><circle cx="12" cy="12" r="4"/>
          <line x1="3" y1="9" x2="6" y2="9"/><line x1="3" y1="15" x2="6" y2="15"/>
          <line x1="18" y1="9" x2="21" y2="9"/><line x1="18" y1="15" x2="21" y2="15"/>
        </svg>
        画像として保存
      </button>
    </div>

    <div class="card">
      <div class="card-title">トークン一覧</div>
      <table class="token-table">
        <thead>
          <tr><th>#</th><th>Token ID</th><th>Piece</th></tr>
        </thead>
        <tbody>
          {token_table_html}
        </tbody>
      </table>
    </div>
    '''}

    {"" if not raw_output else f'''
    <div class="card">
      <details>
        <summary>▸ Raw Output</summary>
        <div class="raw-output">{raw_html}</div>
      </details>
    </div>
    '''}
  </div>
</div>

<script>
function capturePanel() {{
  const panel = document.getElementById('capture-panel');
  if (!panel) return;

  const btn = document.querySelector('.btn-capture');
  const origText = btn.innerHTML;
  btn.textContent = '生成中...';
  btn.disabled = true;

  // スクロール位置をリセットして全体を収める
  window.scrollTo(0, 0);

  html2canvas(panel, {{
    backgroundColor: '#161616',
    scale: 2,
    useCORS: true,
    logging: false,
    windowWidth: panel.scrollWidth + 60,
  }}).then(canvas => {{
    const link = document.createElement('a');
    const modelName = '{cap_model}'.replace(/\.gguf$/i, '') || 'model';
    const ts = new Date().toISOString().slice(0,19).replace(/[:-]/g,'');
    link.download = `tokens_${{modelName}}_${{ts}}.png`;
    link.href = canvas.toDataURL('image/png');
    link.click();
    btn.innerHTML = origText;
    btn.disabled = false;
  }}).catch(err => {{
    console.error(err);
    btn.innerHTML = origText;
    btn.disabled = false;
    alert('キャプチャに失敗しました: ' + err.message);
  }});
}}
</script>
</body>
</html>
"""


class Handler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        print(f"[{self.address_string()}] {format % args}")

    def do_GET(self):
        models = get_gguf_models()
        body = build_html(models).encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", len(body))
        self.end_headers()
        self.wfile.write(body)

    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        raw = self.rfile.read(length).decode("utf-8")
        params = parse_qs(raw)

        model = params.get("model", [""])[0]
        prompt = params.get("prompt", [""])[0]
        models = get_gguf_models()

        tokens = None
        total = None
        raw_output = ""
        error = None

        if not model:
            error = "モデルを選択してください。"
        elif not prompt.strip():
            error = "プロンプトを入力してください。"
        else:
            raw_output, error = run_llama_tokenize(model, prompt)
            if raw_output and not error:
                tokens, total = parse_tokens(raw_output)
                if not tokens:
                    error = "トークンをパースできませんでした。Raw Outputを確認してください。"

        body = build_html(models, model, prompt, raw_output or "", tokens, total, error).encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", len(body))
        self.end_headers()
        self.wfile.write(body)


if __name__ == "__main__":
    print(f"🦙 Tokenizer Visualizer")
    print(f"   llama-tokenize : {LLAMA_TOKENIZE_BIN}")
    print(f"   Models dir     : {GGUF_MODELS_DIR}")
    models = get_gguf_models()
    print(f"   Found models   : {len(models)}")
    for m in models:
        print(f"     - {os.path.basename(m)}")
    print(f"\n   http://localhost:{PORT}/  を開いてください\n")
    server = HTTPServer(("0.0.0.0", PORT), Handler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n停止しました。")