Create app.py

app.py

@@ -0,0 +1,1911 @@
+"""
+app.py — SAE Feature Explorer for Qwen3 models, whether pretrain (base) or posttrain (thinking/instruct) models.
+"""
+
+import argparse
+import html as _html
+import json as _json
+import os
+from collections import OrderedDict
+
+import gradio as gr
+import spaces
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# ─── CLI arguments ────────────────────────────────────────────────────────────
+_parser = argparse.ArgumentParser(description="SAE Feature Explorer")
+_parser.add_argument(
+    '--model',
+    default='Qwen/Qwen3.5-2B',
+    help='Path to the base model directory (default: %(default)s)',
+)
+_parser.add_argument(
+    '--model-name-sae-trained-from',
+    default='qwen3.5-2b-base',
+    help='The name of model which present representations for SAE training (default: %(default)s)',
+)
+_parser.add_argument(
+    '--model-name-analyzing-now',
+    default='qwen3.5-2b',
+    help='The name of model which is used for analyzing now (default: %(default)s)',
+)
+_parser.add_argument(
+    '--sae-path',
+    default='Qwen/SAE-Res-Qwen3.5-2B-Base-W32K-L0_100',
+    help='Path or HF Hub repo ID to the directory containing layer*.sae.pt files (default: %(default)s)',
+)
+_parser.add_argument(
+    '--top-k',
+    type=int,
+    default=100,
+    help='Number of top features to display (default: %(default)s)',
+)
+_parser.add_argument(
+    '--num-layers',
+    type=int,
+    default=24,
+    help='Number of transformer layers in the model (default: %(default)s)',
+)
+_parser.add_argument(
+    '--sae-width',
+    type=int,
+    default=32768,
+    help='SAE dictionary width / number of features (default: %(default)s)',
+)
+_parser.add_argument(
+    '--d-model',
+    type=int,
+    default=2048,
+    help='Model hidden dimension (default: %(default)s)',
+)
+_parser.add_argument(
+    '--sae-cache-max',
+    type=int,
+    default=8,
+    help='Maximum number of SAE layers to keep in memory at once (default: %(default)s)',
+)
+_parser.add_argument(
+    '--server-port',
+    type=int,
+    default=7860,
+    help='Port number for server',
+)
+_args, _unknown = _parser.parse_known_args()
+
+# ─── Config ──────────────────────────────────────────────────────────────────
+MODEL_PATH                  = _args.model
+MODEL_NAME_SAE_TRAINED_FROM = _args.model_name_sae_trained_from
+MODEL_NAME_ANALYZING_NOW    = _args.model_name_analyzing_now
+SAE_PATH                    = _args.sae_path
+TOP_K                       = _args.top_k
+NUM_LAYERS                  = _args.num_layers
+SAE_WIDTH                   = _args.sae_width
+D_MODEL                     = _args.d_model
+SAE_CACHE_MAX               = _args.sae_cache_max
+PORT                        = _args.server_port
+
+# ─── Generation defaults (from model's generation_config.json) ────────────────
+
+_gen_cfg: dict = {}
+_gen_cfg_path = os.path.join(MODEL_PATH, 'generation_config.json')
+if os.path.exists(_gen_cfg_path):
+    with open(_gen_cfg_path) as _f:
+        _gen_cfg = _json.load(_f)
+    print(f"Loaded generation_config.json from {_gen_cfg_path}")
+else:
+    print(f"No generation_config.json found at {_gen_cfg_path}; using built-in defaults.")
+
+GEN_DO_SAMPLE    = bool(_gen_cfg.get('do_sample',           False))
+GEN_TEMPERATURE  = float(_gen_cfg.get('temperature',        1.0))
+GEN_TOP_P        = float(_gen_cfg.get('top_p',              1.0))
+GEN_TOP_K        = int(_gen_cfg.get('top_k',                1))
+GEN_REP_PENALTY  = float(_gen_cfg.get('repetition_penalty', 1.0))
+STEER_DISPLAY_K  = 10   # top-k candidates shown in the per-token probability panel
+
+# ─── Default chat templates (thinking / no-thinking) ─────────────────────────
+
+_THINK_TEMPLATE = (
+    "<|im_start|>user\n"
+    "{content}"
+    "<|im_end|>\n"
+    "<|im_start|>assistant\n"
+    "<think>\n"
+)
+
+_NOTHINK_TEMPLATE = (
+    "<|im_start|>user\n"
+    "{content}"
+    "<|im_end|>\n"
+    "<|im_start|>assistant\n"
+    "<think>\n\n</think>\n\n"
+)
+
+def apply_default_template(prompt: str, think: bool) -> str:
+    """Wrap *prompt* in the ChatML template for thinking or no-thinking mode."""
+    tpl = _THINK_TEMPLATE if think else _NOTHINK_TEMPLATE
+    return tpl.format(content=prompt.strip())
+
+# ─── Device resolution ───────────────────────────────────────────────────────
+
+def _resolve_sae_device() -> torch.device:
+    """
+    Pick the device for SAE weights and encoder/decoder computations.
+
+    CUDA_VISIBLE_DEVICES remaps physical GPUs so that the first listed GPU
+    always appears as cuda:0 inside this process.  We simply use cuda:0
+    when any CUDA device is visible; fall back to CPU otherwise.
+    """
+    if not torch.cuda.is_available():
+        print("SAE device: cpu  (no CUDA visible)")
+        return torch.device('cpu')
+    cvd = os.environ.get('CUDA_VISIBLE_DEVICES', '<unset>')
+    device = torch.device('cuda:0')
+    print(f"SAE device: {device} — {torch.cuda.get_device_name(device)}"
+          f"  [CUDA_VISIBLE_DEVICES={cvd}]")
+    return device
+
+SAE_DEVICE = _resolve_sae_device()
+
+# ─── Global singletons ───────────────────────────────────────────────────────
+_model     = None
+_tokenizer = None
+_sae_lru: OrderedDict = OrderedDict()
+_sae_local_dir: str | None = None   # cached local dir for HF Hub downloaded SAEs
+_orig_cache: dict | None = None   # cached unsteered generation result
+
+
+@spaces.GPU(duration=120)
+def get_model():
+    global _model, _tokenizer
+    if _model is None:
+        print("Loading model…")
+        _model = AutoModelForCausalLM.from_pretrained(
+            MODEL_PATH, device_map='auto', torch_dtype='auto'
+        )
+        _tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+        _model.eval()
+        print("Model ready.")
+    return _model, _tokenizer
+
+
+def get_sae(layer: int) -> dict:
+    if layer in _sae_lru:
+        _sae_lru.move_to_end(layer)
+        return _sae_lru[layer]
+    if len(_sae_lru) >= SAE_CACHE_MAX:
+        _sae_lru.popitem(last=False)
+    # Support both local paths and HF Hub repo IDs
+    global _sae_local_dir
+    if os.path.isdir(SAE_PATH):
+        path = os.path.join(SAE_PATH, f'layer{layer}.sae.pt')
+    else:
+        # Assume HF Hub repo ID – download once, then read from local cache.
+        if _sae_local_dir is None:
+            _sae_local_dir = snapshot_download(SAE_PATH, cache_dir='./sae_cache', local_files_only=False)
+        path = os.path.join(_sae_local_dir, f'layer{layer}.sae.pt')
+    try:
+        sae = torch.load(path, map_location=SAE_DEVICE, weights_only=True)
+    except TypeError:
+        sae = torch.load(path, map_location=SAE_DEVICE)
+    # Pre-convert and transpose encoder weights once on load so compute_sae_features
+    # never repeats the conversion on every call.
+    sae['_W_enc'] = sae['W_enc'].T.to(dtype=torch.float32)  # [d_model, sae_width]
+    sae['_b_enc'] = sae['b_enc'].to(dtype=torch.float32)      # [sae_width]
+    _sae_lru[layer] = sae
+    return sae
+
+
+# ─── Core math ───────────────────────────────────────────────────────────────
+
+def topk_relu(x: torch.Tensor, k: int = TOP_K) -> torch.Tensor:
+    # Scatter top-k ReLU values directly — avoids creating a full-size boolean mask
+    # and an element-wise multiply, saving two [seq, SAE_WIDTH] allocations.
+    relu_x = torch.relu(x)
+    values, indices = torch.topk(relu_x, k, dim=-1)
+    out = torch.zeros_like(relu_x)
+    out.scatter_(-1, indices, values)
+    return out
+
+
+@torch.no_grad()
+def capture_hidden(model, input_ids: torch.Tensor, layer: int) -> torch.Tensor:
+    buf = {}
+    def _hook(module, inp, out):
+        # Qwen3MoE decoder layers return a plain tensor [batch, seq, hidden].
+        # out[0] removes the batch dim → [seq, hidden]; then move to SAE_DEVICE.
+        buf['h'] = out[0].detach().to(SAE_DEVICE, dtype=torch.float32)
+    handle = model.model.layers[layer].register_forward_hook(_hook)
+    model(input_ids)
+    handle.remove()
+    return buf['h']   # [seq_len, d_model]
+
+
+@torch.no_grad()
+def capture_all_hiddens(model, input_ids: torch.Tensor, layers: list) -> dict:
+    """
+    Capture residual-stream hidden states at multiple layers in a single
+    forward pass by registering simultaneous hooks.  Tensors are stored on
+    SAE_DEVICE as float32 so downstream SAE matmuls need no extra transfer.
+    """
+    buf = {}
+    handles = []
+    for layer in layers:
+        def make_hook(l):
+            def _hook(module, inp, out):
+                buf[l] = out[0].detach().to(SAE_DEVICE, dtype=torch.float32)
+            return _hook
+        handles.append(model.model.layers[layer].register_forward_hook(make_hook(layer)))
+    model(input_ids)
+    for h in handles:
+        h.remove()
+    return buf  # {layer_idx: Tensor[seq, d_model] on SAE_DEVICE}
+
+
+def compute_sae_features(hidden: torch.Tensor, sae: dict,
+                         raw: bool = False) -> torch.Tensor:
+    # Use pre-converted weights cached on load (avoids .float()/.T on every call)
+    W_enc = sae['_W_enc']   # [d_model, sae_width] float32 on SAE_DEVICE
+    b_enc = sae['_b_enc']   # [sae_width] float32 on SAE_DEVICE
+    pre   = hidden @ W_enc + b_enc            # [seq, sae_width] — pre-activation on SAE_DEVICE
+    if raw:
+        return pre                            # keep negative values intact; caller handles device
+    return topk_relu(pre, TOP_K)              # stays on SAE_DEVICE; caller calls .tolist() as needed
+
+
+# ─── UI helpers ──────────────────────────────────────────────────────────────
+
+
+def parse_positions(s: str):
+    """
+    Parse a position string into 'all' or a sorted list of int indices.
+
+    Supported syntax (comma-separated, combinable):
+      all        → every token position
+      5          → single position
+      3-7        → inclusive range (positions 3, 4, 5, 6, 7)
+      0,2,5-8    → mix of individual positions and ranges
+    """
+    s = s.strip().lower()
+    if s == 'all':
+        return 'all'
+    try:
+        positions: list[int] = []
+        for part in s.split(','):
+            part = part.strip()
+            if not part:
+                continue
+            if '-' in part:
+                lo, hi = part.split('-', 1)
+                positions.extend(range(int(lo.strip()), int(hi.strip()) + 1))
+            else:
+                positions.append(int(part))
+        return sorted(set(positions))
+    except Exception:
+        return 'all'
+
+
+def feature_heatmap_to_html(tokens: list, features: torch.Tensor, top_k: int, skip_first: bool = False) -> str:
+    """
+    Build a 2-D HTML heatmap:
+      rows  = top-k features (ranked by mean activation across all positions)
+      cols  = token positions
+      color = activation value (white → red, normalised per feature row by row max)
+    """
+
+    seq_len, sae_width = features.shape
+    top_k = min(int(top_k), sae_width)
+
+    # ── Optionally exclude the first token ────────────────────────────────────
+    if skip_first and seq_len > 1:
+        features = features[1:]
+        tokens   = tokens[1:]
+        seq_len -= 1
+
+    # ── Select top-k features by mean activation across all positions ─────────
+    mean_per_feat = features.mean(dim=0)                # [sae_width]
+    top_vals, top_idx = torch.topk(mean_per_feat, top_k)
+    feat_acts = features[:, top_idx]                    # [seq_len, top_k]
+
+    # ── Token column headers ──────────────────────────────────────────────────
+    TH_STYLE = (
+        "min-width:38px;max-width:70px;padding:4px 3px;"
+        "text-align:center;font-weight:500;font-size:11px;"
+        "color:#444;border-bottom:2px solid #c7d2e8;"
+        "overflow:hidden;white-space:nowrap;vertical-align:bottom;"
+    )
+    tok_headers = []
+    for i, tok in enumerate(tokens):
+        raw   = tok.strip() or f"[{i}]"
+        short = _html.escape(raw[:6] + "…" if len(raw) > 6 else raw)
+        full  = _html.escape(raw)
+        tok_headers.append(
+            f'<th style="{TH_STYLE}" title="pos {i}: {full}">{short}</th>'
+        )
+
+    # ── Data rows ─────────────────────────────────────────────────────────────
+    FEAT_TD  = (
+        "font-family:ui-monospace,monospace;font-size:11px;"
+        "padding:3px 8px;color:#2563eb;white-space:nowrap;"
+        "border-right:2px solid #c7d2e8;background:#f8faff;"
+        "position:sticky;left:0;z-index:1;"
+    )
+    AVG_TD   = (
+        "font-size:10px;padding:3px 6px;color:#777;white-space:nowrap;"
+        "border-right:1px solid #e4e7ef;text-align:right;"
+    )
+    CELL_BASE = (
+        "border:1px solid rgba(0,0,0,0.05);min-width:38px;height:30px;"
+        "text-align:center;vertical-align:middle;"
+    )
+
+    rows_html = []
+    for fi in range(top_k):
+        feat_i  = int(top_idx[fi])
+        avg_val = float(top_vals[fi])
+        row_acts = feat_acts[:, fi]             # [seq_len]
+        row_max  = float(row_acts.max())
+        norm     = row_max if row_max > 0 else 1.0
+
+        cells = []
+        for pos in range(seq_len):
+            v = float(row_acts[pos])
+            t = max(0.0, min(1.0, v / norm))
+            # white → amber → deep red
+            r = 255
+            g = int(255 * (1 - 0.8 * t))
+            b = int(255 * (1 - t))
+            cells.append(
+                f'<td style="{CELL_BASE}background:rgb({r},{g},{b});"'
+                f' title="feat #{feat_i} | pos {pos} | act={v:.4f}">'
+                f'</td>'
+            )
+
+        rows_html.append(
+            f'<tr>'
+            f'<td style="{FEAT_TD}">#{feat_i}</td>'
+            f'<td style="{AVG_TD}">{avg_val:.3f}</td>'
+            + "".join(cells)
+            + "</tr>"
+        )
+
+    # ── Assemble table ────────────────────────────────────────────────────────
+    header_row = (
+        '<tr>'
+        '<th style="padding:4px 8px;text-align:left;font-size:11px;font-weight:700;'
+        'color:#2563eb;border-bottom:2px solid #c7d2e8;border-right:2px solid #c7d2e8;'
+        'background:#f8faff;position:sticky;left:0;z-index:2;">Feature</th>'
+        '<th style="padding:4px 6px;font-size:11px;font-weight:700;color:#777;'
+        'border-bottom:2px solid #c7d2e8;border-right:1px solid #e4e7ef;">'
+        'Avg&nbsp;act.</th>'
+        + "".join(tok_headers)
+        + "</tr>"
+    )
+
+    legend = (
+        '<div style="display:flex;align-items:center;gap:10px;margin-top:10px;'
+        'font-size:11px;color:#888;">'
+        '<span>0</span>'
+        '<div style="width:140px;height:12px;border-radius:6px;'
+        'background:linear-gradient(to right,#fff,#ff6600,#cc0000);'
+        'border:1px solid #ddd;"></div>'
+        '<span>peak activation (per-feature row-max scale)</span>'
+        '</div>'
+    )
+
+    return (
+        '<div style="overflow-x:auto;max-height:520px;overflow-y:auto;">'
+        '<table style="border-collapse:collapse;width:100%;'
+        'font-family:ui-monospace,monospace;">'
+        f'<thead style="position:sticky;top:0;background:#fff;z-index:3;">'
+        f'{header_row}</thead>'
+        f'<tbody>{"".join(rows_html)}</tbody>'
+        '</table>'
+        '</div>'
+        + legend
+    )
+
+
+def tokens_with_positions_html(tokens: list, positions) -> str:
+    """
+    Render tokenized prompt as coloured token chips.
+
+    Steered positions (amber/gold) are visually distinct from unsteered ones (grey).
+    positions: 'all'  →  every index is highlighted
+               list   →  only those indices
+    """
+
+    if not tokens:
+        return (
+            '<div style="padding:10px;color:#bbb;font-size:13px;">'
+            'Enter a prompt above to preview token positions.</div>'
+        )
+
+    all_positions = positions if isinstance(positions, list) else []
+    pos_set = (
+        set(range(len(tokens))) if positions == 'all'
+        else {p for p in all_positions if 0 <= p < len(tokens)}
+    )
+    # Positions beyond the prompt — will be steered in the generated text
+    generated_positions = (
+        [] if positions == 'all'
+        else sorted(p for p in all_positions if p >= len(tokens))
+    )
+
+    parts = []
+    for i, tok in enumerate(tokens):
+        steered = i in pos_set
+        txt   = _html.escape(tok)
+        title = _html.escape(repr(tok.strip()), quote=True)
+
+        if steered:
+            bg, border, text_color = "#fef3c7", "2px solid #f59e0b", "#92400e"
+        else:
+            bg, border, text_color = "#f1f5f9", "1px solid #e2e8f0", "#475569"
+
+        parts.append(
+            f'<span style="background:{bg};color:{text_color};'
+            f'padding:3px 7px;margin:2px 1px;border-radius:5px;'
+            f'display:inline-block;border:{border};'
+            f'font-family:ui-monospace,monospace;font-size:12px;" '
+            f'title="pos {i}: {title}">'
+            f'<sub style="opacity:.55;font-size:9px;margin-right:2px">{i}</sub>'
+            f'{txt}</span>'
+        )
+
+    n_steered = len(pos_set)
+    summary = (
+        f'<div style="margin-top:6px;font-size:11px;color:#888;">'
+        f'{len(tokens)}&nbsp;tokens total&nbsp;&nbsp;·&nbsp;&nbsp;'
+        f'<span style="color:#92400e;font-weight:600;">{n_steered}&nbsp;steered</span>'
+        f'&nbsp;<span style="color:#f59e0b;">■</span>'
+        f'</div>'
+    )
+
+    generated_note = ''
+    if generated_positions:
+        gp_str = ', '.join(str(p) for p in generated_positions)
+        generated_note = (
+            f'<div style="margin-top:4px;font-size:11px;padding:4px 8px;'
+            f'background:#eff6ff;border:1px solid #bfdbfe;border-radius:4px;color:#1d4ed8;">'
+            f'Positions {gp_str} are beyond the prompt — they will be steered '
+            f'in the <em>generated</em> text during autoregressive decoding.'
+            f'</div>'
+        )
+
+    return (
+        '<div style="padding:8px 4px;line-height:2.8;">'
+        + ' '.join(parts)
+        + summary
+        + generated_note
+        + '</div>'
+    )
+
+
+def cb_feature_heatmap(state, top_k: int, skip_first: bool):
+    if state is None:
+        return (
+            '<div style="min-height:80px;display:flex;align-items:center;'
+            'justify-content:center;color:#bbb;font-size:13px;">'
+            'Run analysis first to see the feature heatmap.</div>'
+        )
+    tokens, features = state
+    return feature_heatmap_to_html(tokens, features, int(top_k), bool(skip_first))
+
+
+# ─── Gradio callbacks ────────────────────────────────────────────────────────
+
+@spaces.GPU(duration=120)
+def cb_analyze(text: str, layer: int):
+    try:
+        model, tokenizer = get_model()
+        input_ids = tokenizer.encode(text, return_tensors='pt').to(
+            next(model.parameters()).device
+        )
+        tokens   = [tokenizer.decode([t]) for t in input_ids[0].tolist()]
+        hidden   = capture_hidden(model, input_ids, int(layer))
+        features = compute_sae_features(hidden, get_sae(int(layer)))
+        return (tokens, features)
+    except Exception as e:
+        raise gr.Error(f"Analysis failed: {e}")
+
+
+
+def _steering_strength_from_mode(mode: str, diff_lookup, layer: int, feat_idx: int,
+                                  custom_val: float = 5.0) -> float:
+    """Map Light/Medium/Strong/Custom to an actual steering strength.
+
+    Looks up the feature-specific diff for (layer, feat_idx) from the
+    Feature Comparison results.  Falls back to the global max across all
+    compared features, then to fixed defaults when no data is available.
+    """
+    if mode == "Custom":
+        return float(custom_val)
+    d = 0.0
+    if diff_lookup and isinstance(diff_lookup, dict):
+        key = (int(layer), int(feat_idx))
+        if key in diff_lookup:
+            d = float(diff_lookup[key])
+        else:
+            d = float(max(diff_lookup.values(), default=0.0))
+    if d <= 0:
+        return {"Light": 5.0, "Medium": 20.0, "Strong": 100.0}.get(mode, 5.0)
+    return {"Light": round(d * 0.5, 2),
+            "Medium": round(d * 2.0, 2),
+            "Strong": round(d * 10.0, 2)}.get(mode, round(d, 2))
+
+
+@spaces.GPU(duration=120)
+def cb_generate(prompt, layer, feat_idx, pos_str, steer_mode, compare_diff,
+                steer_output_only, max_tok, greedy, top_k_tok, top_p, rep_penalty, temp,
+                custom_strength=5.0, apply_think=False, apply_nothink=False):
+    try:
+        return _cb_generate_inner(prompt, layer, feat_idx, pos_str, steer_mode, compare_diff,
+                                  steer_output_only, max_tok, greedy, top_k_tok, top_p, rep_penalty, temp,
+                                  custom_strength, apply_think, apply_nothink)
+    except gr.Error:
+        raise
+    except Exception as e:
+        raise gr.Error(f"Generation failed: {e}")
+
+
+def cb_update_steer_preview(prompt: str, pos_str: str,
+                             apply_think: bool = False, apply_nothink: bool = False):
+    """Tokenise the prompt and return an HTML token-position preview."""
+    if not prompt.strip():
+        return (
+            '<div style="padding:10px;color:#bbb;font-size:13px;">'
+            'Enter a prompt above to preview steered positions.</div>'
+        )
+    try:
+        _, tokenizer = get_model()
+        if apply_think:
+            effective = apply_default_template(prompt, think=True)
+        elif apply_nothink:
+            effective = apply_default_template(prompt, think=False)
+        else:
+            effective = prompt
+        input_ids = tokenizer.encode(effective)
+        tokens    = [tokenizer.decode([t]) for t in input_ids]
+        positions = parse_positions(pos_str)
+        return tokens_with_positions_html(tokens, positions)
+    except Exception as e:
+        return (
+            f'<div style="padding:10px;color:#dc2626;font-size:13px;">'
+            f'Preview error: {e}</div>'
+        )
+
+
+def _extract_probs(gen_out, input_len: int, tokenizer, display_k: int):
+    """
+    Extract per-step token probabilities from a `return_dict_in_generate=True,
+    output_scores=True` GenerateOutput.
+
+    Returns (text, tokens, chosen_probs, topk_data) where:
+      tokens       : list[str]  — decoded token strings
+      chosen_probs : list[float] — probability of the chosen token (0-1)
+      topk_data    : list[list[[str, float, bool]]] — top-k candidates at each step,
+                     each entry is [token_str, prob, is_chosen]
+    """
+    new_ids     = gen_out.sequences[0][input_len:]
+    new_id_list = new_ids.tolist()
+
+    # Batch-decode chosen tokens and all top-k candidates in two passes
+    # instead of O(n * display_k) individual tokenizer.decode() calls.
+    all_topk_ids: list[list[int]] = []
+    chosen_probs: list[float]     = []
+    topk_vals_list: list          = []
+    chosen_in_top_list: list[bool]= []
+
+    for score_t, tok_id in zip(gen_out.scores, new_id_list):
+        probs              = torch.softmax(score_t[0].float(), dim=-1)
+        chosen_probs.append(float(probs[tok_id]))
+        top_vals, top_ids  = torch.topk(probs, display_k)
+        tid_list           = top_ids.tolist()
+        chosen_in_top      = tok_id in tid_list
+        all_topk_ids.append(tid_list)
+        topk_vals_list.append(top_vals.tolist())
+        chosen_in_top_list.append(chosen_in_top)
+
+    # Single batch_decode call for all chosen tokens
+    tokens: list[str] = tokenizer.batch_decode(
+        [[t] for t in new_id_list], skip_special_tokens=False
+    )
+
+    # Single batch_decode call for all top-k candidate tokens
+    flat_ids      = [tid for ids in all_topk_ids for tid in ids]
+    flat_decoded  = tokenizer.batch_decode(
+        [[t] for t in flat_ids], skip_special_tokens=False
+    )
+
+    topk_data = []
+    flat_idx  = 0
+    for i, (tok_id, ids, vals, chosen_in_top, chosen_prob) in enumerate(
+        zip(new_id_list, all_topk_ids, topk_vals_list, chosen_in_top_list, chosen_probs)
+    ):
+        entry = []
+        for tid, tv in zip(ids, vals):
+            entry.append([flat_decoded[flat_idx], tv, tid == tok_id])
+            flat_idx += 1
+        if not chosen_in_top:
+            entry.append([tokens[i], chosen_prob, True])
+        topk_data.append(entry)
+
+    text = tokenizer.decode(new_ids, skip_special_tokens=True)
+    return text, tokens, chosen_probs, topk_data
+
+
+def probs_to_html(tokens: list, chosen_probs: list, topk_data: list,
+                  panel_id: str, theme: str = 'blue') -> str:
+    """
+    Render per-token generation probabilities as coloured chips.
+    Clicking a chip pins its top-k candidate table in the panel below;
+    clicking the same chip again or another chip toggles/switches the display.
+    Scroll-stable: no hover events that fire on page scroll.
+
+    theme: 'blue' for original output, 'red' for steered output.
+    """
+
+    if not tokens:
+        return ('<div style="padding:10px;color:#bbb;font-size:13px;">'
+                'No tokens generated.</div>')
+
+    # ── Chip colour (white → saturated) based on probability ─────────────────
+    def _colors(prob: float):
+        t = max(0.0, min(1.0, prob))
+        if theme == 'blue':
+            r, g, b = int(255 * (1 - t * 0.85)), int(255 * (1 - t * 0.65)), 255
+            txt = '#1e3a8a' if t < 0.55 else '#ffffff'
+        else:
+            r, g, b = 255, int(255 * (1 - t * 0.82)), int(255 * (1 - t))
+            txt = '#7f1d1d' if t < 0.55 else '#ffffff'
+        return f'rgb({r},{g},{b})', txt
+
+    # ── Pre-build the top-k panel HTML in Python ──────────────────────────────
+    TH = 'padding:2px 8px;font-size:11px;color:#6b7280;border-bottom:1px solid #e4e7ef;'
+
+    def _panel_html(entry: list) -> str:
+        rows = []
+        for rank, (tok_str, prob, is_chosen) in enumerate(entry, 1):
+            bg = 'background:#dbeafe;' if is_chosen else ''
+            fw = 'font-weight:700;'    if is_chosen else ''
+            mk = ' ✓'                  if is_chosen else ''
+            rows.append(
+                f'<tr style="border-bottom:1px solid #f4f6ff;{bg}">'
+                f'<td style="padding:2px 8px;text-align:right;font-size:11px;color:#9ca3af;">{rank}</td>'
+                f'<td style="padding:2px 8px;font-family:monospace;font-size:12px;{fw}">{_html.escape(tok_str)}{mk}</td>'
+                f'<td style="padding:2px 8px;text-align:right;font-family:monospace;font-size:12px;">{prob:.4f}</td>'
+                f'<td style="padding:2px 8px;text-align:right;font-family:monospace;font-size:12px;">{prob * 100:.2f}%</td>'
+                f'</tr>'
+            )
+        return (
+            '<table style="border-collapse:collapse;width:100%;font-size:12px;">'
+            f'<thead style="background:#f8faff;"><tr>'
+            f'<th style="{TH}text-align:right;">Rank</th>'
+            f'<th style="{TH}text-align:left;">Token</th>'
+            f'<th style="{TH}text-align:right;">Prob</th>'
+            f'<th style="{TH}text-align:right;">%</th>'
+            f'</tr></thead>'
+            f'<tbody>{"".join(rows)}</tbody>'
+            '</table>'
+        )
+
+    # ── Inline JS — click to pin, click again to unpin ───────────────────────
+    # Uses data-prob-root to scope sibling chips without global IDs.
+    # Single-quoted JS string literals are safe inside double-quoted HTML attrs.
+    # Non-f-string parts: { } are literal characters (no f-string substitution).
+    JS_CLICK = (
+        "var root=this.closest('[data-prob-root]');"
+        "if(!root)return;"
+        "var p=root.querySelector('[data-topk-panel]');"
+        "if(!p)return;"
+        "var sel=this.dataset.selected==='1';"
+        "root.querySelectorAll('[data-chip]').forEach(function(e){"
+        "e.dataset.selected='0';e.style.outline='';});"
+        "if(sel){"
+        "p.innerHTML='';p.style.display='none';"
+        "}else{"
+        "this.dataset.selected='1';"
+        "this.style.outline='2px solid #94a3b8';"
+        "this.style.outlineOffset='-1px';"
+        "p.innerHTML=this.getAttribute('data-panel');"
+        "p.style.display='block';"
+        "}"
+    )
+
+    def _tok_disp(s: str) -> str:
+        return s.replace('\n', '↵').replace('\r', '↵').replace('\t', '→')
+
+    # ── Build chips ───────────────────────────────────────────────────────────
+    chips = []
+    for tok, prob, entry in zip(tokens, chosen_probs, topk_data):
+        bg, txt = _colors(prob)
+        panel_attr = _html.escape(_panel_html(entry), quote=True)
+        chips.append(
+            f'<span data-chip data-selected="0" '
+            f'style="background:{bg};color:{txt};padding:3px 8px 2px;margin:1px;'
+            f'border-radius:5px;display:inline-block;cursor:pointer;white-space:nowrap;'
+            f'font-family:ui-monospace,monospace;font-size:12px;" '
+            f'data-panel="{panel_attr}" '
+            f'onclick="{JS_CLICK}">'
+            f'{_html.escape(_tok_disp(tok))}'
+            f'<sub style="opacity:.75;font-size:9px;margin-left:3px;">{prob * 100:.1f}%</sub>'
+            f'</span>'
+        )
+
+    return (
+        '<div data-prob-root style="padding:2px;">'
+        '<div style="font-size:11px;color:#888;margin-bottom:6px;font-style:italic;">'
+        'Click a token to pin its top-k candidates &nbsp;·&nbsp; click again to dismiss.</div>'
+        '<div style="padding:4px;line-height:2.8;">'
+        + ''.join(chips)
+        + '</div>'
+        + '<div data-topk-panel style="display:none;margin-top:8px;padding:4px;'
+          'background:#f8faff;border:1px solid #e4e7ef;border-radius:6px;'
+          'max-height:220px;overflow-y:auto;"></div>'
+        + '</div>'
+    )
+
+
+def _cb_generate_inner(prompt, layer, feat_idx, pos_str, steer_mode, compare_diff,
+                       steer_output_only, max_tok, greedy, top_k_tok, top_p, rep_penalty, temp,
+                       custom_strength=5.0, apply_think=False, apply_nothink=False):
+    global _orig_cache
+    model, tokenizer = get_model()
+    layer    = int(layer)
+    feat_idx = int(feat_idx)
+    if not (0 <= feat_idx < SAE_WIDTH):
+        raise gr.Error(f"Feature index must be in [0, {SAE_WIDTH - 1}].")
+    strength  = _steering_strength_from_mode(steer_mode, compare_diff, layer, feat_idx, float(custom_strength))
+    positions = parse_positions(pos_str)
+
+    if apply_think:
+        effective_prompt = apply_default_template(prompt, think=True)
+    elif apply_nothink:
+        effective_prompt = apply_default_template(prompt, think=False)
+    else:
+        effective_prompt = prompt
+    input_ids = tokenizer.encode(effective_prompt, return_tensors='pt').to(
+        next(model.parameters()).device
+    )
+
+    # Build generation kwargs shared by both calls
+    gen_kwargs: dict = dict(max_new_tokens=int(max_tok),
+                            return_dict_in_generate=True, output_scores=True)
+    if greedy:
+        gen_kwargs['do_sample'] = False
+    else:
+        gen_kwargs['do_sample']              = True
+        gen_kwargs['temperature']            = float(temp)
+        gen_kwargs['top_k']                  = int(top_k_tok)
+        gen_kwargs['top_p']                  = float(top_p)
+        gen_kwargs['repetition_penalty']     = float(rep_penalty)
+
+    prompt_len = input_ids.shape[1]
+
+    # ── Original generation (cached) ─────────────────────────────────────────
+    # The unsteered output depends only on the prompt and decoding parameters,
+    # not on any steering inputs.  Reuse the last result when those are unchanged.
+    if greedy:
+        orig_key = (effective_prompt, int(max_tok), True)
+    else:
+        orig_key = (effective_prompt, int(max_tok), False,
+                    int(top_k_tok), float(top_p), float(rep_penalty), float(temp))
+
+    if _orig_cache is not None and _orig_cache['key'] == orig_key:
+        orig_text       = _orig_cache['text']
+        orig_probs_html = _orig_cache['probs_html']
+    else:
+        with torch.no_grad():
+            orig_out = model.generate(input_ids, **gen_kwargs)
+        orig_text, orig_toks, orig_probs, orig_topk = _extract_probs(
+            orig_out, prompt_len, tokenizer, STEER_DISPLAY_K
+        )
+        orig_probs_html = probs_to_html(orig_toks, orig_probs, orig_topk,
+                                        'topk-panel-orig', theme='blue')
+        _orig_cache = dict(key=orig_key, text=orig_text, probs_html=orig_probs_html)
+
+    sae          = get_sae(layer)
+    steering_vec = sae['W_dec'][:, feat_idx].float()   # [d_model]
+    pos_set      = None if positions == 'all' else set(positions)
+    counter      = [0]
+
+    def _steer_hook(module, inp, out):
+        # out: plain tensor [batch, seq, hidden] for Qwen3MoE
+        h           = out.clone()
+        sv          = steering_vec.to(device=h.device, dtype=h.dtype)  # one fused transfer
+        cur_counter = counter[0]
+        counter[0] += 1
+        if cur_counter == 0:
+            # Prefill: apply position-based steering to the prompt
+            if positions == 'all':
+                h = h + strength * sv
+            else:
+                for p in positions:
+                    if 0 <= p < h.shape[1]:
+                        h[:, p, :] = h[:, p, :] + strength * sv
+        else:
+            # Decode step (KV-cache): h is [batch, 1, hidden]
+            # Steer if: output-only mode is on, positions='all', or this position is listed
+            cur_seq_pos = prompt_len + cur_counter - 1
+            if steer_output_only or positions == 'all' or cur_seq_pos in pos_set:
+                h[:, 0, :] = h[:, 0, :] + strength * sv
+        return h
+
+    handle = model.model.layers[layer].register_forward_hook(_steer_hook)
+    with torch.no_grad():
+        steer_out = model.generate(input_ids, **gen_kwargs)
+    handle.remove()
+    steer_text, steer_toks, steer_probs, steer_topk = _extract_probs(
+        steer_out, prompt_len, tokenizer, STEER_DISPLAY_K
+    )
+
+    steer_probs_html = probs_to_html(steer_toks, steer_probs, steer_topk,
+                                     'topk-panel-steer', theme='red')
+
+    return orig_text, steer_text, orig_probs_html, steer_probs_html
+
+
+
+# ─── Feature Comparison helpers ──────────────────────────────────────────────
+
+def compare_to_html(records: list, text1: str, text2: str,
+                    tokens1: list = None, tokens2: list = None) -> tuple:
+    """
+    Render comparison results as two HTML strings:
+      - tok_display_html: token rows for the left panel (data-tok-display root)
+      - feature_table_html: feature table for the right panel
+
+    Returns (tok_display_html, feature_table_html).
+    """
+
+    _TOK_PLACEHOLDER = (
+        '<div style="min-height:60px;display:flex;align-items:center;'
+        'justify-content:center;color:#bbb;font-size:13px;padding:8px;">'
+        'Run Compare to see token activations here.</div>'
+    )
+
+    if not records:
+        return (
+            _TOK_PLACEHOLDER,
+            '<div style="min-height:80px;display:flex;align-items:center;'
+            'justify-content:center;color:#bbb;font-size:13px;">'
+            'No results — try a wider layer range or larger Top-K.</div>',
+        )
+
+    # ── Token display blocks ──────────────────────────────────────────────────
+    TOK_SPAN = (
+        "display:inline-block;padding:3px 7px;margin:2px 1px;"
+        "border-radius:5px;font-family:ui-monospace,monospace;font-size:12px;"
+        "background:#eef2ff;color:#374151;cursor:default;"
+        "transition:background .1s;border:1px solid rgba(0,0,0,0.06);"
+    )
+
+    def render_tok_row(tokens, seq_id):
+        parts = []
+        for i, tok in enumerate(tokens):
+            txt   = _html.escape(tok)
+            title = _html.escape(repr(tok.strip()), quote=True)
+            parts.append(
+                f'<span data-seq={seq_id} data-pos={i} style="{TOK_SPAN}" '
+                f'title="pos {i}: {title}">'
+                f'<sub style="opacity:.5;font-size:9px;margin-right:2px">{i}</sub>'
+                f'{txt}</span>'
+            )
+        return " ".join(parts)
+
+    # Build token display HTML for the left panel
+    if tokens1 and tokens2:
+        tok_inner = (
+            '<div style="margin-bottom:10px;color:#1e293b;">'
+            '<div style="font-size:11px;font-weight:700;color:#2563eb;'
+            'text-transform:uppercase;letter-spacing:.5px;margin-bottom:5px;">'
+            f'Example 1 &nbsp;<span style="font-weight:400;color:#888;">'
+            f'({len(tokens1)} tokens)</span></div>'
+            '<div style="line-height:2.8;padding:8px 10px;background:#fafbff;'
+            'border-radius:8px;border:1px solid #e4e7ef;overflow-x:auto;">'
+            + render_tok_row(tokens1, 1)
+            + '</div></div>'
+            '<div style="margin-bottom:8px;color:#1e293b;">'
+            '<div style="font-size:11px;font-weight:700;color:#dc2626;'
+            'text-transform:uppercase;letter-spacing:.5px;margin-bottom:5px;">'
+            f'Example 2 &nbsp;<span style="font-weight:400;color:#888;">'
+            f'({len(tokens2)} tokens)</span></div>'
+            '<div style="line-height:2.8;padding:8px 10px;background:#fafbff;'
+            'border-radius:8px;border:1px solid #e4e7ef;overflow-x:auto;">'
+            + render_tok_row(tokens2, 2)
+            + '</div></div>'
+            '<div style="font-size:11px;color:#888;font-style:italic;">'
+            'Hover a feature row on the right to highlight activations.</div>'
+        )
+    else:
+        tok_inner = _TOK_PLACEHOLDER
+
+    # Wrap with data-tok-display so the JS hover handler can find it across columns
+    tok_display_html = f'<div data-tok-display style="padding:2px;">{tok_inner}</div>'
+
+    # ── Per-layer max for bar-width normalization ─────────────────────────────
+    _layer_max: dict = {}
+    for _d, _l, *_ in records:
+        if _d > _layer_max.get(_l, 0.0):
+            _layer_max[_l] = _d
+
+    # ── Inline JS snippets for hover-highlight ────────────────────────────────
+    # Uses document.querySelector('[data-tok-display]') so the handler works
+    # across Gradio columns (token panel on left, feature table on right).
+    _JS_ENTER = (
+        "var d=document.querySelector('[data-tok-display]');"
+        "if(!d)return;"
+        "var a1=JSON.parse(this.getAttribute('data-acts1'));"
+        "var a2=JSON.parse(this.getAttribute('data-acts2'));"
+        "if(!a1||!a2)return;"
+        "var pk=Math.max.apply(null,a1.map(Math.abs).concat(a2.map(Math.abs)))||0.0001;"
+        "function c1(v){var t=Math.abs(v)/pk;"
+        "return 'rgb('+Math.round(255*(1-t))+','+Math.round(255*(1-.6*t))+',255)'}"
+        "function c2(v){var t=Math.abs(v)/pk;"
+        "return 'rgb(255,'+Math.round(255*(1-.8*t))+','+Math.round(255*(1-t))+')'}"
+        "d.querySelectorAll('[data-seq]').forEach(function(e){"
+        "var s=e.dataset.seq,p=parseInt(e.dataset.pos,10);"
+        "if(s==='1'&&p<a1.length)e.style.background=c1(a1[p]);"
+        "else if(s==='2'&&p<a2.length)e.style.background=c2(a2[p]);});"
+        "this.style.outline='2px solid #94a3b8';"
+        "this.style.outlineOffset='-1px';"
+    )
+    _JS_LEAVE = (
+        "var d=document.querySelector('[data-tok-display]');"
+        "if(!d)return;"
+        "d.querySelectorAll('[data-seq]').forEach(function(e){e.style.background='';});"
+        "this.style.outline='';"
+    )
+
+    TR_BASE = "border-bottom:1px solid #f0f4ff;"
+    TH = (
+        "padding:6px 10px;font-size:11px;font-weight:700;text-transform:uppercase;"
+        "letter-spacing:.5px;white-space:nowrap;"
+    )
+
+    rows_html = []
+    current_layer = None
+    layer_rank = 0
+    for _rank, record in enumerate(records, 1):
+        diff_val, layer, feat_idx, act1, act2 = record[:5]
+        acts1_pos = record[5] if len(record) > 5 else None
+        acts2_pos = record[6] if len(record) > 6 else None
+
+        # Insert a layer-group header row whenever the layer changes
+        if layer != current_layer:
+            current_layer = layer
+            layer_rank = 0
+            rows_html.append(
+                f'<tr style="background:#eef2ff;border-top:2px solid #c7d2e8;">'
+                f'<td colspan="6" style="padding:4px 12px;font-size:11px;font-weight:700;'
+                f'color:#2563eb;letter-spacing:.5px;">Layer {layer}</td>'
+                f'</tr>'
+            )
+        layer_rank += 1
+
+        bar_w = max(2, int(120 * diff_val / (_layer_max.get(layer) or 1.0)))
+        if act1 >= act2:
+            bar_color = "#2563eb"
+            dir_label = "Ex&nbsp;1&nbsp;▲"
+            dir_color = "#2563eb"
+            row_bg    = "background:#f5f8ff;"
+        else:
+            bar_color = "#dc2626"
+            dir_label = "Ex&nbsp;2&nbsp;▲"
+            dir_color = "#dc2626"
+            row_bg    = "background:#fff5f5;"
+
+        # Embed per-position activation arrays for the hover handler
+        if acts1_pos is not None and acts2_pos is not None:
+            a1_json = _json.dumps(acts1_pos)
+            a2_json = _json.dumps(acts2_pos)
+            tr_open = (
+                f"<tr style='{TR_BASE}{row_bg}cursor:pointer;'"
+                f" data-acts1='{a1_json}'"
+                f" data-acts2='{a2_json}'"
+                f' onmouseenter="{_JS_ENTER}"'
+                f' onmouseleave="{_JS_LEAVE}">'
+            )
+        else:
+            tr_open = f'<tr style="{TR_BASE}{row_bg}">'
+
+        rows_html.append(
+            tr_open
+            + f'<td style="padding:5px 10px;text-align:center;color:#9ca3af;font-size:11px;">{layer_rank}</td>'
+            + f'<td style="padding:5px 10px;font-family:monospace;color:#2563eb;">#{feat_idx}</td>'
+            + f'<td style="padding:5px 8px;text-align:right;font-family:monospace;color:#374151;">{act1:.1%}</td>'
+            + f'<td style="padding:5px 8px;text-align:right;font-family:monospace;color:#374151;">{act2:.1%}</td>'
+            + f'<td style="padding:5px 10px;">'
+            +  f'  <div style="display:flex;align-items:center;gap:6px;">'
+            +  f'    <div style="width:{bar_w}px;height:10px;background:{bar_color};'
+            +  f'         border-radius:3px;flex-shrink:0;"></div>'
+            +  f'    <span style="font-family:monospace;font-size:12px;color:#374151;">{diff_val:.1%}</span>'
+            +  f'  </div>'
+            + f'</td>'
+            + f'<td style="padding:5px 10px;font-size:11px;font-weight:700;color:{dir_color};">'
+            + f'{dir_label}</td>'
+            + '</tr>'
+        )
+
+    ex1_short = _html.escape(text1[:50] + "…" if len(text1) > 50 else text1)
+    ex2_short = _html.escape(text2[:50] + "…" if len(text2) > 50 else text2)
+
+    legend = (
+        '<div style="display:flex;flex-wrap:wrap;gap:16px;margin-top:12px;'
+        'font-size:11px;color:#6b7280;">'
+        f'<span><span style="color:#2563eb;font-weight:700;">■ Ex 1</span>'
+        f' "{ex1_short}"</span>'
+        f'<span><span style="color:#dc2626;font-weight:700;">■ Ex 2</span>'
+        f' "{ex2_short}"</span>'
+        '</div>'
+    )
+
+    table_inner = (
+        '<div style="overflow-x:auto;max-height:560px;overflow-y:auto;color:#1e293b;">'
+        '<table style="border-collapse:collapse;width:100%;color:#1e293b;'
+        'font-family:ui-monospace,monospace;font-size:13px;">'
+        '<thead style="background:#f8faff;color:#1e293b;border-bottom:2px solid #c7d2e8;'
+        'position:sticky;top:0;z-index:2;">'
+        '<tr>'
+        f'<th style="{TH}color:#9ca3af;">Rank</th>'
+        f'<th style="{TH}color:#2563eb;">Feature</th>'
+        f'<th style="{TH}color:#2563eb;text-align:right;">Rate&nbsp;Ex&nbsp;1</th>'
+        f'<th style="{TH}color:#dc2626;text-align:right;">Rate&nbsp;Ex&nbsp;2</th>'
+        f'<th style="{TH}color:#6b7280;">|Rate diff|</th>'
+        f'<th style="{TH}color:#6b7280;">Higher</th>'
+        '</tr>'
+        '</thead>'
+        f'<tbody>{"".join(rows_html)}</tbody>'
+        '</table>'
+        '</div>'
+    )
+
+    feature_table_html = (
+        '<div style="padding:2px;">'
+        + table_inner
+        + legend
+        + '</div>'
+    )
+
+    return tok_display_html, feature_table_html
+
+
+@spaces.GPU(duration=180)
+def cb_compare(text1: str, text2: str, layer_from: int, layer_to: int,
+               top_k: int, skip_first: bool,
+               remove_common_toks: bool, remove_prefix: bool,
+               raw_acts: bool = False):
+    try:
+        if not text1.strip() or not text2.strip():
+            raise gr.Error("Both examples must be non-empty.")
+
+        model, tokenizer = get_model()
+        layer_from = int(layer_from)
+        layer_to   = int(layer_to)
+        top_k      = int(top_k)
+        if layer_from > layer_to:
+            layer_from, layer_to = layer_to, layer_from
+        layers = list(range(layer_from, layer_to + 1))
+
+        # ── Tokenise ─────────────────────────────────────────────────────────
+        model_dev = next(model.parameters()).device
+        ids1 = tokenizer.encode(text1, return_tensors='pt').to(model_dev)
+        ids2 = tokenizer.encode(text2, return_tensors='pt').to(model_dev)
+        toks1 = ids1[0].tolist()
+        toks2 = ids2[0].tolist()
+
+        # ── Build per-sequence keep-index lists ───────────────────────────────
+        prefix_len = 0
+        if remove_prefix:
+            for a, b in zip(toks1, toks2):
+                if a == b:
+                    prefix_len += 1
+                else:
+                    break
+
+        common_tok_ids: set = set()
+        if remove_common_toks:
+            common_tok_ids = set(toks1) & set(toks2)
+
+        def _build_keep(toks: list) -> list:
+            return [
+                i for i, t in enumerate(toks)
+                if not (skip_first and i == 0)
+                and i >= prefix_len
+                and t not in common_tok_ids
+            ]
+
+        keep1 = _build_keep(toks1)
+        keep2 = _build_keep(toks2)
+
+        # ── Capture hidden states for all layers in two forward passes ────────
+        hiddens1 = capture_all_hiddens(model, ids1, layers)
+        hiddens2 = capture_all_hiddens(model, ids2, layers)
+
+        # Decoded token strings for the HTML token display
+        tokens1_str = [tokenizer.decode([t]) for t in toks1]
+        tokens2_str = [tokenizer.decode([t]) for t in toks2]
+
+        # ── Per-layer feature activation-rate difference ──────────────────────
+        # Activation rate = fraction of kept positions where the feature fires
+        # (activation > 0).  Ranking by |rate1 − rate2| highlights features
+        # that are selectively active in one example but not the other.
+        # Load one SAE at a time to avoid OOM (each SAE is ~1-2 GB on GPU).
+        candidates = []   # (abs_diff, layer, feat_idx, rate1, rate2,
+                          #  acts1_per_pos, acts2_per_pos)
+        for layer in layers:
+            sae = get_sae(layer)
+
+            # Full per-position feature activations — stay on SAE_DEVICE for GPU math
+            feats1 = compute_sae_features(hiddens1[layer], sae, raw=raw_acts)  # [seq1_len, SAE_WIDTH]
+            feats2 = compute_sae_features(hiddens2[layer], sae, raw=raw_acts)  # [seq2_len, SAE_WIDTH]
+
+            # Activation rate = fraction of kept positions where feature fires (> 0)
+            def _rate(feats: torch.Tensor, keep_idx: list) -> torch.Tensor:
+                if not keep_idx:
+                    return torch.zeros(feats.shape[1], device=feats.device, dtype=feats.dtype)
+                return (feats[keep_idx] > 0).float().mean(dim=0)
+
+            r1   = _rate(feats1, keep1)
+            r2   = _rate(feats2, keep2)
+            diff = (r1 - r2).abs()
+
+            # Top-k per layer (all kept — no global trim)
+            local_k    = min(top_k, SAE_WIDTH)
+            vals, idxs = torch.topk(diff, local_k)
+            for v, fi in zip(vals.tolist(), idxs.tolist()):
+                # Round to 3 dp — enough precision for color interpolation
+                a1_pos = [round(x, 3) for x in feats1[:, fi].tolist()]
+                a2_pos = [round(x, 3) for x in feats2[:, fi].tolist()]
+                candidates.append((v, layer, fi, float(r1[fi]), float(r2[fi]),
+                                   a1_pos, a2_pos))
+
+            # Free SAE weights and feature tensors before loading the next layer
+            del sae, feats1, feats2, diff
+
+        # Single cache clear after all layers — calling it per-layer is expensive
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        # ── Per-layer sort: group by layer, within each layer sort by diff desc ─
+        candidates.sort(key=lambda x: (x[1], -x[0]))
+        diff_lookup: dict = {}
+        for diff_val, layer, feat_idx, *_ in candidates:
+            key = (layer, feat_idx)
+            if key not in diff_lookup or diff_val > diff_lookup[key]:
+                diff_lookup[key] = diff_val
+        tok_html, table_html = compare_to_html(candidates, text1, text2, tokens1_str, tokens2_str)
+        return tok_html, table_html, diff_lookup
+
+    except gr.Error:
+        raise
+    except Exception as e:
+        raise gr.Error(f"Comparison failed: {e}")
+
+
+# ─── CSS ─────────────────────────────────────────────────────────────────────
+
+CSS = """
+/* ══════════════════════════════════════════════════════════════════
+   Color tokens — single source of truth for light / dark palettes
+   ══════════════════════════════════════════════════════════════════ */
+:root {
+    --c-page-bg:      #f4f6fb;
+    --c-card-bg:      #ffffff;
+    --c-card-border:  #e4e7ef;
+    --c-card-shadow:  0 1px 4px rgba(0,0,0,0.06), 0 4px 16px rgba(0,0,0,0.04);
+    --c-header-bg:    linear-gradient(135deg,#eff6ff 0%,#e0eaff 55%,#ede9fe 100%);
+    --c-header-border:#c7d2fe;
+    --c-header-text:  #1e293b;
+    --c-header-h1:    #1e3a8a;
+    --c-header-p:     #475569;
+    --c-pill-bg:      rgba(37,99,235,0.08);
+    --c-pill-border:  rgba(37,99,235,0.22);
+    --c-pill-text:    #1e3a8a;
+    --c-chip-bg:      #eff4ff;
+    --c-chip-text:    #2563eb;
+    --c-btn2-bg:      #f8faff;
+    --c-btn2-border:  #d0d7e8;
+    --c-btn2-text:    #374151;
+    --c-outbox-bg:    #fafbff;
+    --c-outbox-text:  #1e293b;
+    --c-outbox-border:#e4e7ef;
+    --c-tab-text:     #374151;
+    --c-tab-sel:      #2563eb;
+    --c-divider:      #dde3f0;
+    --c-th-bg:        #f0f4ff;
+    --c-th-text:      #2563eb;
+}
+
+/* Dark mode via OS/browser preference */
+@media (prefers-color-scheme: dark) {
+    :root {
+        --c-page-bg:      #0f172a;
+        --c-card-bg:      #1e293b;
+        --c-card-border:  #334155;
+        --c-card-shadow:  0 1px 4px rgba(0,0,0,0.40), 0 4px 16px rgba(0,0,0,0.25);
+        --c-header-bg:    linear-gradient(135deg,#172554 0%,#1e3a8a 55%,#3b0764 100%);
+        --c-header-border:#1e40af;
+        --c-header-text:  #e2e8f0;
+        --c-header-h1:    #bfdbfe;
+        --c-header-p:     #94a3b8;
+        --c-pill-bg:      rgba(96,165,250,0.12);
+        --c-pill-border:  rgba(96,165,250,0.30);
+        --c-pill-text:    #93c5fd;
+        --c-chip-bg:      #172554;
+        --c-chip-text:    #93c5fd;
+        --c-btn2-bg:      #1e293b;
+        --c-btn2-border:  #475569;
+        --c-btn2-text:    #e2e8f0;
+        --c-outbox-bg:    #0f172a;
+        --c-outbox-text:  #e2e8f0;
+        --c-outbox-border:#334155;
+        --c-tab-text:     #94a3b8;
+        --c-tab-sel:      #60a5fa;
+        --c-divider:      #334155;
+        --c-th-bg:        #172554;
+        --c-th-text:      #93c5fd;
+    }
+}
+
+/* Dark mode via Gradio's explicit dark-mode class (toggled manually) */
+.dark {
+    --c-page-bg:      #0f172a;
+    --c-card-bg:      #1e293b;
+    --c-card-border:  #334155;
+    --c-card-shadow:  0 1px 4px rgba(0,0,0,0.40), 0 4px 16px rgba(0,0,0,0.25);
+    --c-header-bg:    linear-gradient(135deg,#172554 0%,#1e3a8a 55%,#3b0764 100%);
+    --c-header-border:#1e40af;
+    --c-header-text:  #e2e8f0;
+    --c-header-h1:    #bfdbfe;
+    --c-header-p:     #94a3b8;
+    --c-pill-bg:      rgba(96,165,250,0.12);
+    --c-pill-border:  rgba(96,165,250,0.30);
+    --c-pill-text:    #93c5fd;
+    --c-chip-bg:      #172554;
+    --c-chip-text:    #93c5fd;
+    --c-btn2-bg:      #1e293b;
+    --c-btn2-border:  #475569;
+    --c-btn2-text:    #e2e8f0;
+    --c-outbox-bg:    #0f172a;
+    --c-outbox-text:  #e2e8f0;
+    --c-outbox-border:#334155;
+    --c-tab-text:     #94a3b8;
+    --c-tab-sel:      #60a5fa;
+    --c-divider:      #334155;
+    --c-th-bg:        #172554;
+    --c-th-text:      #93c5fd;
+}
+
+/* ── Page background ── */
+body, .gradio-container { background: var(--c-page-bg) !important; }
+
+/* ── Header card ── */
+.header-card {
+    background: var(--c-header-bg);
+    border-radius: 14px;
+    padding: 22px 28px 18px;
+    margin-bottom: 4px;
+    color: var(--c-header-text);
+    box-shadow: 0 4px 20px rgba(37,99,235,0.10);
+    border: 1px solid var(--c-header-border);
+}
+.header-card h1 { margin:0 0 6px; font-size:24px; font-weight:700; letter-spacing:-.3px; color:var(--c-header-h1); }
+.header-card p  { margin:0;       font-size:13px; color:var(--c-header-p); }
+.stat-pill {
+    display:inline-block;
+    background:var(--c-pill-bg);
+    border:1px solid var(--c-pill-border);
+    border-radius:20px;
+    padding:3px 13px;
+    font-size:12px;
+    color:var(--c-pill-text);
+    margin:4px 3px 0;
+}
+
+/* ── Panel cards ── */
+.panel-card {
+    background: var(--c-card-bg) !important;
+    border-radius: 12px !important;
+    box-shadow: var(--c-card-shadow) !important;
+    border: 1px solid var(--c-card-border) !important;
+    padding: 18px !important;
+}
+.panel-card > .form { gap: 12px !important; }
+
+/* ── Section label chips ── */
+.section-chip {
+    font-size: 11px;
+    font-weight: 700;
+    text-transform: uppercase;
+    letter-spacing: .8px;
+    color: var(--c-chip-text);
+    background: var(--c-chip-bg);
+    border-radius: 6px;
+    padding: 2px 10px;
+    display: inline-block;
+    margin-bottom: 10px;
+}
+
+/* ── Buttons ── */
+.btn-primary {
+    background: linear-gradient(135deg, #2563eb, #6d28d9) !important;
+    border: none !important;
+    border-radius: 8px !important;
+    font-weight: 600 !important;
+    font-size: 14px !important;
+    letter-spacing: .2px !important;
+    box-shadow: 0 2px 10px rgba(37,99,235,0.30) !important;
+    transition: all 0.18s ease !important;
+    color: #fff !important;
+    padding: 10px 0 !important;
+}
+.btn-primary:hover {
+    transform: translateY(-1px) !important;
+    box-shadow: 0 5px 18px rgba(37,99,235,0.40) !important;
+}
+.btn-secondary {
+    border-radius: 8px !important;
+    font-weight: 500 !important;
+    font-size: 13px !important;
+    border: 1px solid var(--c-btn2-border) !important;
+    background: var(--c-btn2-bg) !important;
+    color: var(--c-btn2-text) !important;
+    transition: all 0.15s ease !important;
+}
+.btn-secondary:hover {
+    background: var(--c-chip-bg) !important;
+    border-color: var(--c-tab-sel) !important;
+}
+
+/* ── Output boxes ── */
+.output-box textarea {
+    font-family: ui-monospace, monospace !important;
+    font-size: 13px !important;
+    line-height: 1.7 !important;
+    background: var(--c-outbox-bg) !important;
+    color: var(--c-outbox-text) !important;
+    border-color: var(--c-outbox-border) !important;
+    border-radius: 8px !important;
+}
+
+/* ── Dataframe ── */
+.feature-table table { font-family: ui-monospace, monospace; font-size: 13px; }
+.feature-table th    { background: var(--c-th-bg) !important; color: var(--c-th-text) !important;
+                        font-weight: 600; font-size: 12px; text-transform: uppercase; }
+
+/* ── Tab styling ── */
+.tab-nav button {
+    font-weight: 600 !important;
+    font-size: 14px !important;
+    border-radius: 8px 8px 0 0 !important;
+    color: var(--c-tab-text) !important;
+}
+.tab-nav button.selected {
+    color: var(--c-tab-sel) !important;
+    border-bottom: 2px solid var(--c-tab-sel) !important;
+}
+
+/* ── Divider ── */
+.section-divider {
+    border: none;
+    border-top: 1px dashed var(--c-divider);
+    margin: 6px 0 10px;
+}
+
+/* ── Slider label ── */
+label.svelte-1b6s6sv { font-size: 13px !important; font-weight: 500 !important; }
+"""
+
+# ─── Build the Gradio interface ───────────────────────────────────────────────
+
+with gr.Blocks(title="Qwen-Scope Feature Explorer", theme=gr.themes.Soft(), css=CSS) as demo:
+
+    # ── Header ────────────────────────────────────────────────────────────────
+    gr.HTML(
+        '<div class="header-card">'
+        '  <div style="display:flex;align-items:center;gap:8px;margin-bottom:6px;">'
+        '    <img src="https://cdn-avatars.huggingface.co/v1/production/uploads/620760a26e3b7210c2ff1943/-s1gyJfvbE1RgO5iBeNOi.png" alt="Qwen Logo" style="height:24px;width:auto;">'
+        '    <h1 style="margin:0;">Qwen-Scope Feature Explorer</h1>'
+        '  </div>'
+        f'  <p>Interpret {MODEL_NAME_ANALYZING_NOW} via Sparse Autoencoders trained on each residual-stream layer from {MODEL_NAME_SAE_TRAINED_FROM}.</p>'
+        '  <div style="margin-top:10px;">'
+        f'    <span class="stat-pill">Model: {MODEL_NAME_ANALYZING_NOW}</span>'
+        f'    <span class="stat-pill">SAE trained from: {MODEL_NAME_SAE_TRAINED_FROM}</span>'
+        f'    <span class="stat-pill">Layers: {NUM_LAYERS}</span>'
+        f'    <span class="stat-pill">SAE width: {SAE_WIDTH:,}</span>'
+        f'    <span class="stat-pill">Top-k: {TOP_K}</span>'
+        f'    <span class="stat-pill">d_model: {D_MODEL}</span>'
+        '  </div>'
+        '</div>'
+    )
+
+    analysis_state = gr.State(None)   # (list[str] tokens, Tensor[seq, sae_width] features)
+    compare_diff_state = gr.State({})
+
+    with gr.Tabs(elem_classes="tab-nav"):
+
+        # ═════════════════════════════════════════════��════════════════════════
+        # Tab 1 — Feature Comparison
+        # ══════════════════════════════════════════════════════════════════════
+        with gr.Tab("⚖️  Feature Comparison"):
+
+            with gr.Row(equal_height=False):
+
+                # ── Left column: inputs + settings + token preview ─────────────
+                with gr.Column(scale=2, min_width=300):
+
+                    with gr.Accordion("Examples", open=True) as t3_examples_accordion:
+                        with gr.Group(elem_classes="panel-card"):
+                            gr.HTML('<span class="section-chip">Examples</span>')
+                            t3_text1 = gr.Textbox(
+                                label="Example 1",
+                                lines=5,
+                                placeholder="Paste first text here…",
+                            )
+                            t3_text2 = gr.Textbox(
+                                label="Example 2",
+                                lines=5,
+                                placeholder="Paste second text here…",
+                            )
+
+                    with gr.Accordion("Comparison Settings", open=True) as t3_settings_accordion:
+                        with gr.Group(elem_classes="panel-card"):
+                            gr.HTML('<span class="section-chip">Comparison Settings</span>')
+                            with gr.Row():
+                                t3_layer_from = gr.Slider(
+                                    minimum=0, maximum=NUM_LAYERS - 1,
+                                    value=0, step=1,
+                                    label="Layer from",
+                                    scale=1,
+                                )
+                                t3_layer_to = gr.Slider(
+                                    minimum=0, maximum=NUM_LAYERS - 1,
+                                    value=NUM_LAYERS - 1, step=1,
+                                    label="Layer to",
+                                    scale=1,
+                                )
+                            t3_topk = gr.Number(
+                                value=5, precision=0,
+                                label="Top-K results",
+                                info="Number of (layer, feature) pairs to surface.",
+                            )
+                            with gr.Accordion("Advanced options", open=False):
+                                t3_skip_first = gr.Checkbox(
+                                    label="Exclude first token",
+                                    value=False,
+                                    info="Skip position 0 when computing mean activations.",
+                                )
+                                t3_remove_common_toks = gr.Checkbox(
+                                    label="Remove common tokens",
+                                    value=False,
+                                    info="Exclude positions whose token ID appears in both examples.",
+                                )
+                                t3_remove_prefix = gr.Checkbox(
+                                    label="Remove common prefix",
+                                    value=False,
+                                    info="Exclude the longest token-level prefix shared by both examples.",
+                                )
+                            t3_run = gr.Button(
+                                "⚖️  Compare Features",
+                                variant="primary",
+                                elem_classes="btn-primary",
+                            )
+
+                    with gr.Accordion("Features", open=True) as t3_features_accordion:
+                        with gr.Group(elem_classes="panel-card"):
+                            gr.HTML(
+                                '<span class="section-chip">Feature Comparison</span>'
+                                '<span style="font-size:12px;color:#888;margin-left:8px;">'
+                                'top-K features per layer · ranked by |rate(Ex1) − rate(Ex2)|'
+                                ' where rate = fraction of token positions where the feature fires · grouped by layer'
+                                '</span>'
+                            )
+                            t3_out = gr.HTML(
+                                value=(
+                                    '<div style="min-height:80px;display:flex;align-items:center;'
+                                    'justify-content:center;color:#bbb;font-size:13px;">'
+                                    'Enter two examples and click Compare.</div>'
+                                )
+                            )
+
+                # ── Right column: token activations ──────────────────────────���
+                with gr.Column(scale=3, min_width=380):
+                    with gr.Group(elem_classes="panel-card"):
+                        gr.HTML(
+                            '<span class="section-chip">Token Activations</span>'
+                            '<span style="font-size:12px;color:#888;margin-left:8px;">'
+                            'hover a feature row on the left to highlight activations'
+                            '</span>'
+                        )
+                        t3_tok_html = gr.HTML(
+                            value=(
+                                '<div style="min-height:60px;display:flex;align-items:center;'
+                                'justify-content:center;color:#bbb;font-size:13px;padding:8px;">'
+                                'Run Compare to see token activations here.</div>'
+                            )
+                        )
+
+        # ══════════════════════════════════════════════════════════════════════
+        # Tab 2 — Feature Steering
+        # ══════════════════════════════════════════════════════════════════════
+        with gr.Tab("🎛️  Feature Steering"):
+
+            with gr.Row(equal_height=False):
+
+                # ── Left column: prompt + steering controls ────────────────
+                with gr.Column(scale=2, min_width=280):
+                    with gr.Group(elem_classes="panel-card"):
+                        gr.HTML('<span class="section-chip">Prompt</span>')
+                        t2_prompt = gr.Textbox(
+                            label=None,
+                            lines=5,
+                            placeholder="Enter a generation prompt…",
+                            show_label=False,
+                        )
+
+                        t2_apply_think = gr.Checkbox(
+                            label="Apply default thinking template",
+                            value=False,
+                            info=(
+                                "Wrap the prompt in the ChatML format with thinking enabled "
+                                "(assistant prefill starts with <think>)."
+                            ),
+                        )
+                        t2_apply_nothink = gr.Checkbox(
+                            label="Apply default no-thinking template",
+                            value=False,
+                            info=(
+                                "Wrap the prompt in the ChatML format with thinking disabled "
+                                "(assistant prefill starts with <think>\\n\\n</think>)."
+                            ),
+                        )
+                        t2_template_info = gr.HTML(visible=False, value="")
+
+                        gr.HTML('<span class="section-chip">Token Position Preview</span>'
+                                '<span style="font-size:12px;color:#888;margin-left:8px;">'
+                                'amber = steered &nbsp;·&nbsp; updates as you type'
+                                '</span>')
+
+                        t2_pos_preview = gr.HTML(
+                            value=(
+                                '<div style="padding:10px;color:#bbb;font-size:13px;">'
+                                'Enter a prompt above to preview steered positions.</div>'
+                            )
+                        )
+
+                    with gr.Group(elem_classes="panel-card"):
+                        gr.HTML('<span class="section-chip">Steering Parameters</span>')
+
+                        with gr.Row():
+                            t2_layer = gr.Slider(
+                                minimum=0, maximum=NUM_LAYERS - 1,
+                                value=10, step=1,
+                                label="Layer",
+                                scale=3,
+                            )
+                            t2_feat = gr.Number(
+                                value=0, precision=0,
+                                label="Feature index",
+                                info=f"0 – {SAE_WIDTH - 1}",
+                                scale=2,
+                            )
+
+                        t2_pos = gr.Textbox(
+                            label="Token positions to steer",
+                            value="all",
+                            placeholder="all  |  0,1,5  |  3-7  |  0,2,5-8",
+                            info=(
+                                "all → every token  |  "
+                                "0,1,5 → individual positions  |  "
+                                "3-7 → inclusive range  |  "
+                                "combinations e.g. 0,2,5-8"
+                            ),
+                        )
+                        t2_steer_output_only = gr.Checkbox(
+                            label="Also steer generated tokens",
+                            value=True,
+                            info=(
+                                "When enabled, every generated token is steered in addition to "
+                                "whatever the positions field specifies for the prompt."
+                            ),
+                        )
+
+                        gr.HTML('<span class="section-chip">Steering Strength</span>')
+                        t2_steer_mode = gr.Radio(
+                            choices=["Light", "Medium", "Strong", "Custom"],
+                            value="Light",
+                            label=None,
+                            show_label=False,
+                            info=(
+                                "Calibrated to the most different feature found in "
+                                "Feature Comparison. Run that tab first."
+                            ),
+                        )
+                        t2_custom_strength = gr.Number(
+                            value=5.0,
+                            label="Custom strength",
+                            info="Direct steering magnitude (used when Custom is selected above).",
+                            visible=False,
+                            precision=2,
+                        )
+                        t2_steer_info = gr.HTML(
+                            value=(
+                                '<div style="font-size:11px;color:#888;padding:4px 6px;'
+                                'background:#f8faff;border-radius:5px;">'
+                                'Light ≈ 5.0 · Medium ≈ 20.0 · Strong ≈ 100.0<br>'
+                                '<span style="color:#bbb;">Run Feature Comparison to calibrate.</span>'
+                                '</div>'
+                            )
+                        )
+
+                        gr.HTML('<hr class="section-divider">')
+                        with gr.Accordion("Sampling options", open=False):
+                            t2_maxtok = gr.Slider(
+                                minimum=20, maximum=300,
+                                value=100, step=10,
+                                label="Max new tokens",
+                            )
+                            t2_greedy = gr.Checkbox(
+                                label="Greedy decoding",
+                                value=True,
+                                info="When enabled, all sampling parameters below are ignored.",
+                            )
+                            with gr.Row():
+                                t2_temperature = gr.Slider(
+                                    minimum=0.01, maximum=2.0,
+                                    value=GEN_TEMPERATURE, step=0.01,
+                                    label="Temperature",
+                                    interactive=GEN_DO_SAMPLE,
+                                )
+                                t2_top_p = gr.Slider(
+                                    minimum=0.0, maximum=1.0,
+                                    value=GEN_TOP_P, step=0.01,
+                                    label="Top-p (nucleus)",
+                                    interactive=GEN_DO_SAMPLE,
+                                )
+                            with gr.Row():
+                                t2_top_k_tok = gr.Slider(
+                                    minimum=0, maximum=200,
+                                    value=GEN_TOP_K, step=1,
+                                    label="Top-k (tokens)",
+                                    info="0 = disabled",
+                                    interactive=GEN_DO_SAMPLE,
+                                )
+                                t2_rep_penalty = gr.Slider(
+                                    minimum=1.0, maximum=3.0,
+                                    value=GEN_REP_PENALTY, step=0.05,
+                                    label="Repetition penalty",
+                                    info="1.0 = no penalty",
+                                    interactive=GEN_DO_SAMPLE,
+                                )
+
+                        t2_run = gr.Button(
+                            "▶  Generate Both Outputs",
+                            variant="primary",
+                            elem_classes="btn-primary",
+                        )
+
+                # ── Right column: outputs ──────────────────────────────────
+                with gr.Column(scale=3, min_width=380):
+
+                    with gr.Group(elem_classes="panel-card"):
+                        gr.HTML(
+                            '<span class="section-chip">Original Output</span>'
+                            '<span style="font-size:12px;color:#888;margin-left:8px;">'
+                            'No steering applied</span>'
+                        )
+                        t2_orig = gr.Textbox(
+                            label=None, lines=7,
+                            interactive=False,
+                            show_label=False,
+                            placeholder="Original generation will appear here…",
+                            elem_classes="output-box",
+                        )
+                        gr.HTML(
+                            '<span class="section-chip" style="margin-top:10px;'
+                            'display:inline-block;">Token Probabilities</span>'
+                            '<span style="font-size:12px;color:#888;margin-left:8px;">'
+                            'blue intensity = confidence &nbsp;·&nbsp; hover = top-k</span>'
+                        )
+                        t2_orig_probs = gr.HTML(
+                            value='<div style="padding:10px;color:#bbb;font-size:13px;">'
+                                  'Run generation to see token probabilities.</div>'
+                        )
+
+                    with gr.Group(elem_classes="panel-card"):
+                        gr.HTML(
+                            '<span class="section-chip" style="background:#fef3f2;color:#dc2626;">'
+                            'Steered Output</span>'
+                            '<span style="font-size:12px;color:#888;margin-left:8px;">'
+                            'With SAE feature injection</span>'
+                        )
+                        t2_steered = gr.Textbox(
+                            label=None, lines=7,
+                            interactive=False,
+                            show_label=False,
+                            placeholder="Steered generation will appear here…",
+                            elem_classes="output-box",
+                        )
+                        gr.HTML(
+                            '<span class="section-chip" style="background:#fef3f2;color:#dc2626;'
+                            'margin-top:10px;display:inline-block;">Token Probabilities</span>'
+                            '<span style="font-size:12px;color:#888;margin-left:8px;">'
+                            'red intensity = confidence &nbsp;·&nbsp; hover = top-k</span>'
+                        )
+                        t2_steer_probs = gr.HTML(
+                            value='<div style="padding:10px;color:#bbb;font-size:13px;">'
+                                  'Run generation to see token probabilities.</div>'
+                        )
+
+            t2_run.click(
+                cb_generate,
+                inputs=[t2_prompt, t2_layer, t2_feat, t2_pos, t2_steer_mode, compare_diff_state,
+                        t2_steer_output_only, t2_maxtok,
+                        t2_greedy, t2_top_k_tok, t2_top_p, t2_rep_penalty,
+                        t2_temperature, t2_custom_strength, t2_apply_think, t2_apply_nothink],
+                outputs=[t2_orig, t2_steered, t2_orig_probs, t2_steer_probs],
+            )
+            t3_run.click(
+                cb_compare,
+                inputs=[t3_text1, t3_text2, t3_layer_from, t3_layer_to, t3_topk,
+                        t3_skip_first, t3_remove_common_toks, t3_remove_prefix],
+                outputs=[t3_tok_html, t3_out, compare_diff_state],
+            ).then(
+                fn=lambda: [gr.update(open=False), gr.update(open=False)],
+                inputs=None,
+                outputs=[t3_examples_accordion, t3_settings_accordion],
+            )
+            _sampling_controls = [
+                t2_temperature, t2_top_p, t2_top_k_tok, t2_rep_penalty
+            ]
+            t2_greedy.change(
+                fn=lambda g: [gr.update(interactive=not g)] * 4,
+                inputs=[t2_greedy],
+                outputs=_sampling_controls,
+            )
+            t2_prompt.change(
+                cb_update_steer_preview,
+                inputs=[t2_prompt, t2_pos, t2_apply_think, t2_apply_nothink],
+                outputs=[t2_pos_preview],
+            )
+            t2_pos.change(
+                cb_update_steer_preview,
+                inputs=[t2_prompt, t2_pos, t2_apply_think, t2_apply_nothink],
+                outputs=[t2_pos_preview],
+            )
+
+            def _update_steer_info(mode: str, diff_lookup, layer, feat_idx):
+                if mode == "Custom":
+                    return (
+                        '<div style="font-size:11px;color:#555;padding:4px 6px;'
+                        'background:#f8faff;border-radius:5px;">'
+                        'Enter a custom steering strength value above.'
+                        '</div>'
+                    )
+                d = 0.0
+                source_note = '<span style="color:#bbb;">Run Feature Comparison to calibrate.</span>'
+                if diff_lookup and isinstance(diff_lookup, dict):
+                    key = (int(layer), int(feat_idx))
+                    if key in diff_lookup:
+                        d = float(diff_lookup[key])
+                        source_note = (
+                            f'<span style="color:#16a34a;">feature #{int(feat_idx)} '
+                            f'@ layer {int(layer)} · diff = {d:.3f}</span>'
+                        )
+                    else:
+                        d = float(max(diff_lookup.values(), default=0.0))
+                        source_note = (
+                            f'<span style="color:#64748b;">feature not in compare results — '
+                            f'using global max diff = {d:.3f}</span>'
+                        )
+                if d <= 0:
+                    vals = {"Light": 5.0, "Medium": 20.0, "Strong": 100.0}
+                else:
+                    vals = {
+                        "Light":  round(d * 0.5, 2),
+                        "Medium": round(d * 2.0, 2),
+                        "Strong": round(d * 10.0, 2),
+                    }
+                return (
+                    f'<div style="font-size:11px;color:#555;padding:4px 6px;'
+                    f'background:#f8faff;border-radius:5px;">'
+                    f'Light ≈ {vals["Light"]} · Medium ≈ {vals["Medium"]} · Strong ≈ {vals["Strong"]}<br>'
+                    + source_note + '</div>'
+                )
+
+            _steer_info_inputs = [t2_steer_mode, compare_diff_state, t2_layer, t2_feat]
+            for _trigger in [t2_steer_mode.change, compare_diff_state.change,
+                             t2_layer.change, t2_feat.change]:
+                _trigger(
+                    fn=_update_steer_info,
+                    inputs=_steer_info_inputs,
+                    outputs=[t2_steer_info],
+                )
+
+            # Show/hide custom strength input depending on radio selection
+            t2_steer_mode.change(
+                fn=lambda m: gr.update(visible=(m == "Custom")),
+                inputs=[t2_steer_mode],
+                outputs=[t2_custom_strength],
+            )
+
+            # ── Template toggle: mutual exclusion + info panel + preview refresh ─
+            _THINK_INFO_HTML = (
+                '<div style="font-size:11px;color:#555;padding:6px 10px;'
+                'background:#eff6ff;border:1px solid #bfdbfe;border-radius:6px;'
+                'font-family:ui-monospace,monospace;white-space:pre-wrap;line-height:1.7;">'
+                '&lt;|im_start|&gt;user\n'
+                '&#123;your prompt&#125;&lt;|im_end|&gt;\n'
+                '&lt;|im_start|&gt;assistant\n'
+                '&lt;think&gt;\n'
+                '</div>'
+            )
+            _NOTHINK_INFO_HTML = (
+                '<div style="font-size:11px;color:#555;padding:6px 10px;'
+                'background:#f0fdf4;border:1px solid #bbf7d0;border-radius:6px;'
+                'font-family:ui-monospace,monospace;white-space:pre-wrap;line-height:1.7;">'
+                '&lt;|im_start|&gt;user\n'
+                '&#123;your prompt&#125;&lt;|im_end|&gt;\n'
+                '&lt;|im_start|&gt;assistant\n'
+                '&lt;think&gt;\n\n&lt;/think&gt;\n\n'
+                '</div>'
+            )
+
+            def _on_think_change(think_val, nothink_val, prompt, pos_str):
+                if think_val:
+                    # Just checked: uncheck nothink, show think format, refresh preview
+                    return (gr.update(value=False),
+                            gr.update(visible=True, value=_THINK_INFO_HTML),
+                            cb_update_steer_preview(prompt, pos_str, True, False))
+                elif nothink_val:
+                    # Unchecked by mutual exclusion — nothink is active; leave info+preview alone
+                    return gr.update(), gr.update(), gr.update()
+                else:
+                    # Manually unchecked with nothing active — reset to raw
+                    return (gr.update(),
+                            gr.update(visible=False),
+                            cb_update_steer_preview(prompt, pos_str, False, False))
+
+            def _on_nothink_change(nothink_val, think_val, prompt, pos_str):
+                if nothink_val:
+                    # Just checked: uncheck think, show nothink format, refresh preview
+                    return (gr.update(value=False),
+                            gr.update(visible=True, value=_NOTHINK_INFO_HTML),
+                            cb_update_steer_preview(prompt, pos_str, False, True))
+                elif think_val:
+                    # Unchecked by mutual exclusion — think is active; leave info+preview alone
+                    return gr.update(), gr.update(), gr.update()
+                else:
+                    # Manually unchecked with nothing active — reset to raw
+                    return (gr.update(),
+                            gr.update(visible=False),
+                            cb_update_steer_preview(prompt, pos_str, False, False))
+
+            t2_apply_think.change(
+                fn=_on_think_change,
+                inputs=[t2_apply_think, t2_apply_nothink, t2_prompt, t2_pos],
+                outputs=[t2_apply_nothink, t2_template_info, t2_pos_preview],
+            )
+            t2_apply_nothink.change(
+                fn=_on_nothink_change,
+                inputs=[t2_apply_nothink, t2_apply_think, t2_prompt, t2_pos],
+                outputs=[t2_apply_think, t2_template_info, t2_pos_preview],
+            )
+
+
+if __name__ == '__main__':
+    # ZeroGPU: model is loaded lazily inside @spaces.GPU decorated functions,
+    # so we do NOT pre-load here (Space boots on CPU, GPU is allocated on demand).
+    print("Starting Gradio server on ZeroGPU Space…")
+    demo.queue(max_size=4)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=int(os.environ.get("PORT", PORT)),
+        share=False,
+        strict_cors=False,
+        show_error=True,
+        ssr_mode=False,
+    )