"""Bertographer — Gradio Space: a mechanistic-interpretability cockpit for encoder *classifiers* (BERT / RoBERTa / DeBERTa / DistilBERT / ELECTRA). The decoder sibling (Cartogemma) reads a causal LM's next-token futures. Bertographer reads an encoder classifier's *decision*: for every layer and attention head, it projects the pooling token's register into TWO spaces — · vocab (logit lens via tied word_embeddings) → what token this register "says" · class (the model's own classification head) → which label this head votes …and lets an operator intervene (mute / scale a head, fire a trigger) and watch the prediction move. Framed for the engine room, not the lab: input = a trace layers = the waterfall heads = spans / voters mute/scale = the counterfactual ("what if the model hadn't seen this evidence?") cmp = trace-vs-trace diff ("what differed between this request and the last?") Default model: cardiffnlp/twitter-roberta-base-sentiment-latest (3-class sentiment). Architecture is auto-discovered; per-head views degrade gracefully when a model's attention output projection isn't a clean per-head concatenation. Methodology lifted from the cartographer family: cartographer/mechanistic_work/negation/deberta_cartographer.py (encoder cell: tok|class|cos) cartographer/mechanistic_work/twitter_roberta/capture_*.py (4 projection spaces, per-head) cartographer/mechanistic_work/twitter_roberta/intervene_sentiment.py (head-slot interventions) gradiographer/cartogemma/app.py (Gradio cockpit + REPL + CSS) """ import os import re import sys import json import unicodedata from dataclasses import dataclass, field import torch import torch.nn as nn import torch.nn.functional as F import gradio as gr from transformers import AutoModelForSequenceClassification, AutoTokenizer, set_seed # ── Config ────────────────────────────────────────────────────────────────── DEFAULT_MODEL_ID = os.environ.get("BERTOGRAPHER_MODEL", "cardiffnlp/twitter-roberta-base-sentiment-latest") HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") SEED = 42 # Resolve device WITHOUT touching CUDA at import time (on ZeroGPU that races the # `spaces` import; on CPU Spaces it's just cpu). Local GPU users: set BERTOGRAPHER_DEVICE=cuda. DEVICE = os.environ.get("BERTOGRAPHER_DEVICE", "cpu") # ── HEAT palette for token-probability (mirrors cartographer3) ────────────── HEAT_HEX = ["#59749C", "#948863", "#ECA60F", "#FFCF67", "#219C7F", "#E0D05B", "#64D32A"] def heat_color(p: float) -> str: if p >= 0.85: return HEAT_HEX[6] if p >= 0.70: return HEAT_HEX[5] if p >= 0.50: return HEAT_HEX[4] if p >= 0.30: return HEAT_HEX[3] if p >= 0.15: return HEAT_HEX[2] if p >= 0.05: return HEAT_HEX[1] return HEAT_HEX[0] # Class colors: semantic for sentiment, palette cycle otherwise. _CLASS_PALETTE = ["#64D32A", "#FFCF67", "#21C5C5", "#E78BFF", "#ff6b6b", "#9aa0ff", "#E0D05B"] def class_color(label_name: str, idx: int) -> str: n = label_name.lower() if n.startswith("pos") or n in ("entailment", "positive", "label_2"): return "#64D32A" if n.startswith("neg") or n in ("contradiction", "negative", "label_0"): return "#ff6b6b" if n.startswith("neu") or n in ("neutral", "label_1"): return "#FFCF67" return _CLASS_PALETTE[idx % len(_CLASS_PALETTE)] def char_width(ch: str) -> int: cat = unicodedata.category(ch) if cat in ("Mn", "Me"): return 0 eaw = unicodedata.east_asian_width(ch) if eaw in ("W", "F"): return 2 return 1 def display_width(s: str) -> int: return sum(char_width(c) for c in s) def fmt_tok(tok: str, max_w: int = 7) -> str: tok = tok.replace("\n", "\\n").replace("\r", "\\r").replace("\t", "\\t") tok = tok.lstrip("▁").lstrip("Ġ") # SP / GPT-2 leading-space markers if not tok: tok = "·" w, out = 0, [] for ch in tok: cw = char_width(ch) if w + cw > max_w: break out.append(ch); w += cw return "".join(out) def html_escape(s: str) -> str: return s.replace("&", "&").replace("<", "<").replace(">", ">") def preprocess_tweet(text: str) -> str: """Cardiff TweetEval normalization — applied only for twitter-* models.""" text = re.sub(r"@\w+", "@user", text) text = re.sub(r"http\S+", "http", text) return text # ── Architecture resolver ─────────────────────────────────────────────────── # An encoder classifier I can fully read needs three things: # 1. an encoder block stack (ModuleList of layers), each with an attention # output projection whose INPUT is the per-head concatenation # (attention.output.dense for BERT/RoBERTa/DeBERTa/ELECTRA, attention.out_lin # for DistilBERT). This gives per-head decomposition (Tier-2). # 2. word_embeddings, for the tied-embedding logit lens (vocab projection). # 3. a classification head I can replay on a single pooled vector. # When (1) doesn't decompose cleanly I degrade to Tier-1 (full-layer only). _BASE_ATTRS = ("roberta", "bert", "deberta", "distilbert", "electra", "xlm_roberta", "camembert", "mpnet", "albert") def _resolve_base(model): for a in _BASE_ATTRS: b = getattr(model, a, None) if b is not None: return b # generic: first child module exposing embeddings + an encoder/transformer for _, mod in model.named_children(): if hasattr(mod, "embeddings") and (hasattr(mod, "encoder") or hasattr(mod, "transformer")): return mod return model def _resolve_layers(base): enc = getattr(base, "encoder", None) if enc is not None and hasattr(enc, "layer") and isinstance(enc.layer, nn.ModuleList): return enc.layer tr = getattr(base, "transformer", None) # DistilBERT if tr is not None and hasattr(tr, "layer") and isinstance(tr.layer, nn.ModuleList): return tr.layer # ALBERT: encoder.albert_layer_groups[0].albert_layers (shared) — best effort if enc is not None and hasattr(enc, "albert_layer_groups"): grp = enc.albert_layer_groups[0] if hasattr(grp, "albert_layers"): return grp.albert_layers # generic: longest attention-bearing ModuleList best = None for _, mod in base.named_modules(): if isinstance(mod, nn.ModuleList) and len(mod) and hasattr(mod[0], "attention"): if best is None or len(mod) > len(best): best = mod if best is None: raise RuntimeError("Could not locate an encoder layer stack for this model.") return best def _resolve_attn_out(layer): """The attention output projection whose input is per-head-concatenated.""" attn = getattr(layer, "attention", None) if attn is not None: out = getattr(attn, "output", None) if out is not None and hasattr(out, "dense"): # BERT/RoBERTa/DeBERTa/ELECTRA return out.dense if hasattr(attn, "out_lin"): # DistilBERT return attn.out_lin if hasattr(attn, "dense"): return attn.dense return None def _resolve_word_embeddings(base, model): emb = getattr(base, "embeddings", None) if emb is not None and hasattr(emb, "word_embeddings"): return emb.word_embeddings ie = model.get_input_embeddings() return ie class ClassHead: """Replays a model's classification head on a single pooled [hidden] vector. Adapts across the five common encoder-classifier head shapes.""" def __init__(self, model): self.kind = None self.fns = [] clf = getattr(model, "classifier", None) cname = type(clf).__name__ if clf is not None else "" if hasattr(model, "pre_classifier") and clf is not None: # DistilBERT: pre_classifier(Linear) -> ReLU -> classifier(Linear) self.kind = "distilbert" pre, c = model.pre_classifier, clf self.fns = [lambda x: torch.relu(pre(x)), c] elif clf is not None and hasattr(clf, "out_proj") and hasattr(clf, "dense"): # RoBERTa(tanh) / ELECTRA(gelu) classification head act = torch.tanh if "electra" in cname.lower() or "electra" in type(model).__name__.lower(): act = F.gelu d, op = clf.dense, clf.out_proj self.kind = "roberta_head" self.fns = [lambda x, d=d, act=act: act(d(x)), op] elif getattr(model, "pooler", None) is not None and clf is not None: # BERT pooler(dense+tanh) or DeBERTa ContextPooler(dense+gelu) -> classifier(Linear) pooler = model.pooler pname = type(pooler).__name__.lower() act = F.gelu if "context" in pname or "deberta" in pname else torch.tanh dense = getattr(pooler, "dense", None) if dense is not None: self.kind = "pooled" self.fns = [lambda x, d=dense, act=act: act(d(x)), clf] else: self.kind = "linear"; self.fns = [clf] elif clf is not None: self.kind = "linear"; self.fns = [clf] else: raise RuntimeError("Could not locate a classification head for this model.") # weight vectors of the final projection, for cosine-to-class final = self.fns[-1] self.final_weight = final.weight.detach() if hasattr(final, "weight") else None def logits(self, vec: torch.Tensor) -> torch.Tensor: x = vec if x.dim() == 1: x = x.unsqueeze(0) for f in self.fns: x = f(x) return x # [1, num_labels] def probs(self, vec: torch.Tensor): with torch.no_grad(): return F.softmax(self.logits(vec).float(), dim=-1)[0] def cosine(self, vec: torch.Tensor, class_idx: int) -> float: if self.final_weight is None or class_idx >= self.final_weight.shape[0]: return 0.0 v = vec.squeeze().float() c = self.final_weight[class_idx].float() return float(F.cosine_similarity(v.unsqueeze(0), c.unsqueeze(0)).item()) @dataclass class Snapshot: """One forward pass, shared by every pane.""" class_logits: torch.Tensor # [num_labels] class_probs: torch.Tensor # [num_labels] hidden_states: tuple # len L+1, each [1, seq, hidden] head_inputs: dict = field(default_factory=dict) # layer -> attn-out input [1, seq, H*hd] pred: int = 0 margin: float = 0.0 entropy: float = 0.0 # over class distribution (bits) # ── The Probe ─────────────────────────────────────────────────────────────── class EncoderProbe: def __init__(self, model_id: str = None): model_id = model_id or DEFAULT_MODEL_ID self.model_id = model_id print(f"[*] Loading {model_id} on {DEVICE}...") kw = {"token": HF_TOKEN} if HF_TOKEN else {} self.tokenizer = AutoTokenizer.from_pretrained(model_id, **kw) self.model = AutoModelForSequenceClassification.from_pretrained(model_id, **kw) self.model.to(DEVICE).eval() self.is_tweet = "twitter" in model_id.lower() cfg = self.model.config self.base = _resolve_base(self.model) self.layers = _resolve_layers(self.base) self.num_layers = len(self.layers) self.num_heads = int(getattr(cfg, "num_attention_heads", 0)) or 1 self.hidden_size = int(getattr(cfg, "hidden_size", 0)) or 0 self.head_dim = self.hidden_size // self.num_heads if self.num_heads else 0 self.word_emb = _resolve_word_embeddings(self.base, self.model) self.vocab_size = self.word_emb.weight.shape[0] self.head = ClassHead(self.model) self.num_labels = int(getattr(cfg, "num_labels", 0)) or self.head.logits( torch.zeros(self.hidden_size, device=DEVICE)).shape[-1] id2label = getattr(cfg, "id2label", None) or {i: f"LABEL_{i}" for i in range(self.num_labels)} self.label_names = [str(id2label.get(i, f"LABEL_{i}")) for i in range(self.num_labels)] self.label_codes = self._make_codes(self.label_names) # Tier-2 (per-head) support: attn-out input must decompose as heads*head_dim. o0 = _resolve_attn_out(self.layers[0]) self._attn_outs = [_resolve_attn_out(l) for l in self.layers] self.head_scan_supported = bool( o0 is not None and self.head_dim and o0.weight.shape[1] == self.num_heads * self.head_dim ) self.arch_name = type(self.model).__name__ print(f"[*] {self.arch_name}: {self.num_layers}L x {self.num_heads}H, " f"head_dim={self.head_dim}, hidden={self.hidden_size}, vocab={self.vocab_size}, " f"labels={self.label_names}, head_scan=" f"{'yes' if self.head_scan_supported else 'NO (Tier-1)'}") self.input_ids = None self.attn_mask = None self.pos = 0 self.snapshot = None self.scaled_heads = {} # (layer, head) -> scale (0.0 == muted) self._iv_handles = [] self.full_log = [] @staticmethod def _make_codes(names): codes, seen = [], set() for n in names: base = re.sub(r"[^A-Za-z0-9]", "", n).upper() or "X" c = base[:3] i = 1 while c in seen: c = (base[:2] + str(i))[:3]; i += 1 seen.add(c); codes.append(c) return codes # ── architecture accessors ── def attn_out(self, li): return self._attn_outs[li] def vocab_logits(self, vec): """Tied-embedding logit lens — encoder classifiers ship no lm_head.""" return F.linear(vec, self.word_emb.weight.to(vec.dtype)) def set_seed(self): set_seed(SEED) def load_input(self, text: str, text_pair: str = None): text = text or "" if self.is_tweet: text = preprocess_tweet(text) if text_pair: text_pair = preprocess_tweet(text_pair) enc = self.tokenizer(text, text_pair, return_tensors="pt", truncation=True, max_length=512) self.input_ids = enc["input_ids"].to(DEVICE) self.attn_mask = enc.get("attention_mask") if self.attn_mask is not None: self.attn_mask = self.attn_mask.to(DEVICE) self.pos = 0 self.full_log.append({"type": "load", "text": text[:200]}) # ── intervention hooks (operate on per-head slots of the attn-out input) ── def _install_iv_hooks(self): for h in self._iv_handles: h.remove() self._iv_handles = [] if not self.scaled_heads or not self.head_scan_supported: return by_layer = {} for (l, h), s in self.scaled_heads.items(): by_layer.setdefault(l, {})[h] = s hd = self.head_dim for li, scales in by_layer.items(): mod = self.attn_out(li) if mod is None: continue def make_hook(scales, hd=hd): def fn(module, args): x = args[0].clone() for h, s in scales.items(): x[:, :, h * hd:(h + 1) * hd] *= s return (x,) + args[1:] return fn self._iv_handles.append(mod.register_forward_pre_hook(make_hook(scales))) def scale_head(self, layer, head, scale): if not self.head_scan_supported: return f"per-head intervention unsupported for {self.arch_name} (Tier-1 only)" if not (0 <= layer < self.num_layers): return f"layer {layer} out of range" if not (0 <= head < self.num_heads): return f"head {head} out of range" if scale == 1.0: self.scaled_heads.pop((layer, head), None) else: self.scaled_heads[(layer, head)] = float(scale) self._install_iv_hooks() self.full_log.append({"type": "scale", "layer": layer, "head": head, "scale": scale}) verb = "MUTED" if scale == 0.0 else (f"x{scale}" if scale != 1.0 else "cleared") return f"L{layer}H{head} {verb}" def clear_interventions(self): n = len(self.scaled_heads) self.scaled_heads.clear() self._install_iv_hooks() self.full_log.append({"type": "clear_iv"}) return f"cleared {n} intervention(s)" # ── single shared forward pass ── def forward_snapshot(self): if self.input_ids is None: self.snapshot = None return None captured = {} handles = [] if self.head_scan_supported: # capture the (already intervention-scaled) input to each attn-out proj for li in range(self.num_layers): op = self.attn_out(li) if op is not None: handles.append(op.register_forward_pre_hook( (lambda li: (lambda m, a: captured.__setitem__(li, a[0].detach())))(li))) try: with torch.no_grad(): out = self.model(self.input_ids, attention_mask=self.attn_mask, output_hidden_states=True) finally: for h in handles: h.remove() logits = out.logits[0].float() probs = F.softmax(logits, dim=-1) sp = torch.sort(probs, descending=True).values margin = float((sp[0] - sp[1]).item()) if probs.numel() > 1 else float(sp[0].item()) ent = float(-(probs * torch.log2(probs + 1e-12)).sum().item()) self.snapshot = Snapshot( class_logits=logits, class_probs=probs, hidden_states=out.hidden_states, head_inputs=captured, pred=int(probs.argmax().item()), margin=margin, entropy=ent, ) return self.snapshot # ── per-head register projections at self.pos ── def _project_vec(self, vec, width=3): """Project a [hidden] register → (top tokens, class probs, top class, cos).""" with torch.no_grad(): vl = self.vocab_logits(vec) vp = F.softmax(vl, dim=-1) tp, ti = torch.topk(vp, width) toks = [(self.tokenizer.decode([ti[i].item()]), float(tp[i].item())) for i in range(width)] cp = self.head.probs(vec) top_c = int(cp.argmax().item()) cos = self.head.cosine(vec, top_c) return {"tokens": toks, "class_probs": cp.cpu(), "top_class": top_c, "cosine": cos} def scan_heads(self, target_layers=None, target_heads=None, width=3): snap = self.snapshot if snap is None: return None if target_layers is None: target_layers = list(range(self.num_layers)) if target_heads is None: target_heads = list(range(self.num_heads)) target_layers = [l for l in target_layers if 0 <= l < self.num_layers] target_heads = [h for h in target_heads if 0 <= h < self.num_heads] pos = min(self.pos, snap.hidden_states[0].shape[1] - 1) rows = {li: {} for li in target_layers} # Layer column: full residual at this position for li in target_layers: vec = snap.hidden_states[li + 1][0, pos, :] rows[li]["Layer"] = self._project_vec(vec, width) if not (self.head_scan_supported and snap.head_inputs): return {"rows": rows, "heads": [], "pos": pos} for li in target_layers: inp = snap.head_inputs.get(li) if inp is None: continue wo = self.attn_out(li).weight # [hidden, hidden] wv = wo.view(wo.shape[0], self.num_heads, self.head_dim) x = inp[0, pos, :].view(self.num_heads, self.head_dim) # [H, hd] # proj[h] = Wo[:, h, :] @ x[h] → [H, hidden] proj_all = torch.einsum("khd,hd->hk", wv.to(x.dtype), x) for hi in target_heads: rows[li][hi] = self._project_vec(proj_all[hi], width) return {"rows": rows, "heads": target_heads, "pos": pos} # ── per-(head,layer) trace of one vocab token's rank ── def vocab_rank_trace(self, token_str, rank_lo=0, rank_hi=None): snap = self.snapshot if snap is None or not (self.head_scan_supported and snap.head_inputs): return "unsupported" if snap is not None else None if token_str.isdigit(): tid = int(token_str); name = self.tokenizer.decode([tid]) else: ids = self.tokenizer.encode(token_str, add_special_tokens=False) if not ids: return None tid = ids[0]; name = token_str if rank_hi is None: rank_hi = self.vocab_size pos = min(self.pos, snap.hidden_states[0].shape[1] - 1) rows, index = [], [] for li in sorted(snap.head_inputs.keys()): wo = self.attn_out(li).weight wv = wo.view(wo.shape[0], self.num_heads, self.head_dim) x = snap.head_inputs[li][0, pos, :].view(self.num_heads, self.head_dim) proj_all = torch.einsum("khd,hd->hk", wv.to(x.dtype), x) for hi in range(self.num_heads): rows.append(proj_all[hi]); index.append((hi, li)) ranks = {} if rows: with torch.no_grad(): logits = self.vocab_logits(torch.stack(rows, 0)) # [N, vocab] target = logits[:, tid].unsqueeze(1) rk = (logits > target).sum(dim=1).add(1).tolist() for (hi, li), r in zip(index, rk): ranks[(hi, li)] = int(r) return {"name": name, "tid": tid, "rank_lo": rank_lo, "rank_hi": rank_hi, "ranks": ranks, "mode": "rank"} # ── per-(head,layer) probability of one class ── def class_prob_trace(self, class_idx): snap = self.snapshot if snap is None or not (self.head_scan_supported and snap.head_inputs): return "unsupported" if snap is not None else None if not (0 <= class_idx < self.num_labels): return None pos = min(self.pos, snap.hidden_states[0].shape[1] - 1) vals = {} for li in sorted(snap.head_inputs.keys()): wo = self.attn_out(li).weight wv = wo.view(wo.shape[0], self.num_heads, self.head_dim) x = snap.head_inputs[li][0, pos, :].view(self.num_heads, self.head_dim) proj_all = torch.einsum("khd,hd->hk", wv.to(x.dtype), x) for hi in range(self.num_heads): vals[(hi, li)] = float(self.head.probs(proj_all[hi])[class_idx].item()) return {"class_idx": class_idx, "name": self.label_names[class_idx], "vals": vals, "mode": "class"} # ── HTML renderers ────────────────────────────────────────────────────────── CELL_W = 6 def render_input(p: EncoderProbe) -> str: if p is None or p.input_ids is None: return "
No trace loaded. Enter text and click Load.
" toks = p.tokenizer.convert_ids_to_tokens(p.input_ids[0]) special = set(p.tokenizer.all_special_tokens) parts = [f"
[{len(toks)} tokens · inspecting pos {p.pos}]
", "
"] for i, t in enumerate(toks): disp = html_escape(t.replace("▁", "·").replace("Ġ", "·")) if i == p.pos: parts.append(f"[{i}:{disp}] ") elif t in special: parts.append(f"[{i}:{disp}] ") else: parts.append(f"{i}:{disp} ") parts.append("
") return "
" + "".join(parts) + "
" def _cell(p, cell): tok, tp = cell["tokens"][0] ci = cell["top_class"] cp = float(cell["class_probs"][ci].item()) code = p.label_codes[ci] tc = heat_color(tp) cc = class_color(p.label_names[ci], ci) return (f"" f"{html_escape(fmt_tok(tok, CELL_W))}" f"{code}{int(cp*100):02d}" f".{int(abs(cell['cosine'])*99):02d}") def render_headmap(p: EncoderProbe, data: dict) -> str: if not data: return ("
No scan yet. h * scans all " "layers, h L one layer.
") rows_data = data["rows"] heads = data["heads"] layers = sorted(rows_data.keys()) if not layers: return "
Empty scan.
" hdr = ["Lay"] for h in heads: hdr.append(f"H{h}") hdr.append("Layer") sub = [""] if heads: sub.append(f"per-head span " f"(head→o_proj→vocab|class)") sub.append("residual") body = [] for l in layers: r = [f"L{l}"] row = rows_data[l] for h in heads: r.append(_cell(p, row[h]) if h in row else "·") r.append(_cell(p, row["Layer"]).replace("hm-cell", "hm-cell full") if "Layer" in row else "·") r.append("") body.append("".join(r)) note = "" if not heads: note = ("
per-head decomposition unavailable " "for this architecture — full-layer residual only (Tier-1)
") return (f"
{note}" f"{''.join(sub)}{''.join(hdr)}{''.join(body)}
" f"
cell: tok (logit-lens) · " f"CODE pp (head's class vote) · .cos (cosine to that class)
") def render_prediction(p: EncoderProbe) -> str: if p is None or p.snapshot is None: return "
Load a trace to see the prediction.
" s = p.snapshot pred = s.pred rows = ["
"] rows.append(f"
prediction " f"" f"{html_escape(p.label_names[pred])} " f"({s.class_probs[pred]*100:.1f}%) · " f"margin {s.margin*100:.1f}pp · H={s.entropy:.2f} bits
") rows.append("") for i in range(p.num_labels): pr = float(s.class_probs[i].item()) c = class_color(p.label_names[i], i) bar = int(round(pr * 28)) mark = " ◀" if i == pred else "" rows.append( f"" f"" f"") rows.append("
{html_escape(p.label_names[i])}{'█'*bar}{'·'*(28-bar)}{pr*100:5.1f}%{mark}
") # final-layer per-head vote tally (the "span attributes" view) if p.head_scan_supported and s.head_inputs: li = p.num_layers - 1 inp = s.head_inputs.get(li) if inp is not None: pos = min(p.pos, s.hidden_states[0].shape[1] - 1) wo = p.attn_out(li).weight wv = wo.view(wo.shape[0], p.num_heads, p.head_dim) x = inp[0, pos, :].view(p.num_heads, p.head_dim) proj_all = torch.einsum("khd,hd->hk", wv.to(x.dtype), x) cells = [] for hi in range(p.num_heads): ci = int(p.head.probs(proj_all[hi]).argmax().item()) c = class_color(p.label_names[ci], ci) cells.append(f"" f"{p.label_codes[ci]}") rows.append(f"
final-layer head votes " f"(L{li}): {' '.join(cells)}
") rows.append("
") return "".join(rows) def render_trace(p: EncoderProbe, td: dict) -> str: if not td: return ("
Trace one signal across heads×layers: " "spark <token> (vocab rank) or " "class <label> (vote probability).
") n_layers, n_heads = p.num_layers, p.num_heads band = 8 head_colors = ["#64D32A", "#FFCF67", "#21C5C5", "#E78BFF", "#ff6b6b", "#9aa0ff"] lines = [] hdr = " " + "".join(f"{l:>3}" for l in range(n_layers)) if td["mode"] == "rank": rank_lo, rank_hi = td["rank_lo"], td["rank_hi"] span = max(rank_hi - rank_lo, 1) ranks = td["ranks"] lines.append(f"vocab rank of '{html_escape(td['name'])}' (id={td['tid']}) — " f"{rank_lo} (top) → {rank_hi} (bottom)") lines.append(f"{html_escape(hdr)}") lines.append(f" {'─'*(n_layers*3)}") for hi in range(n_heads): color = head_colors[hi % len(head_colors)] grid = [[" "]*n_layers for _ in range(band)] for li in range(n_layers): r = ranks.get((hi, li), rank_hi + 1) if r < rank_lo: grid[0][li] = "▲" elif r > rank_hi: grid[band-1][li] = "▼" else: row = int((r - rank_lo)/span*(band-1)) grid[max(0, min(band-1, row))][li] = "●" _emit_band(lines, grid, band, hi, n_layers, color, rank_lo, rank_hi, span) if hi < n_heads - 1: lines.append(f" {'┈'*(n_layers*3)}") lines.append(f" {'─'*(n_layers*3)}") else: vals = td["vals"] c = class_color(td["name"], td["class_idx"]) lines.append(f"P({html_escape(td['name'])}) per head — 1.0 (top) → 0.0 (bottom)") lines.append(f"{html_escape(hdr)}") lines.append(f" {'─'*(n_layers*3)}") for hi in range(n_heads): grid = [[" "]*n_layers for _ in range(band)] for li in range(n_layers): v = vals.get((hi, li), 0.0) row = int((1.0 - v)*(band-1)) grid[max(0, min(band-1, row))][li] = "●" _emit_band(lines, grid, band, hi, n_layers, c, 1.0, 0.0, 1.0, pct=True) if hi < n_heads - 1: lines.append(f" {'┈'*(n_layers*3)}") lines.append(f" {'─'*(n_layers*3)}") return f"
{chr(10).join(lines)}
" def _emit_band(lines, grid, band, hi, n_layers, color, top, bot, span, pct=False): for ri in range(band): if ri == 0: lab = f"H{hi:<2}{(top if not pct else 1.0):>6}│" if not pct else f"H{hi:<2}{1.0:>6.2f}│" elif ri == band - 1: lab = f" {bot:>6}│" if not pct else f" {0.0:>6.2f}│" else: lab = " │" cells = [] for li in range(n_layers): ch = grid[ri][li] if ch in ("●", "▲", "▼"): cells.append(f" {ch} ") else: cells.append(f" {ch} ") lines.append(html_escape(lab) + "".join(cells)) def render_summary(s) -> str: if s.probe is None: return "
no model loaded
" p = s.probe snap = p.snapshot n_tok = p.input_ids.shape[1] if p.input_ids is not None else 0 tier = "T2" if p.head_scan_supported else "T1" pred = (f"" f"{html_escape(p.label_names[snap.pred])} {snap.class_probs[snap.pred]*100:.0f}%" if snap else "—") iv = (f"iv {len(p.scaled_heads)}" if p.scaled_heads else "iv 0") items = [ f"{html_escape(p.arch_name)} {html_escape(p.model_id)}", f"tier {tier}", f"L{p.num_layers} ×H{p.num_heads}", f"tok {n_tok}", f"pos {p.pos}", f"pred {pred}", f"H={snap.entropy:.2f}" if snap else "", iv, ] return "
" + " · ".join(x for x in items if x) + "
" def status_html(s) -> str: if not s.transcript: return "
Ready.
" rows = [] for cmd, st in s.transcript[-22:]: rows.append(f"
› {html_escape(cmd)} " f"— {html_escape(st)}
") return f"
{''.join(rows)}
" # ── Session state ─────────────────────────────────────────────────────────── class Session: def __init__(self): self.probe = None self.scan_data = None self.trace_data = None self.transcript = [] self.cmp = None # (label_names, probs, pred) of a compared trace self.loaded_model = None self.loaded_text = None self.loaded_pair = None def load(self, model_id, text, pair, force=False): changed = (force or self.probe is None or self.loaded_model != model_id or self.loaded_text != text or self.loaded_pair != pair) if not changed: return False if self.probe is None or self.loaded_model != model_id: self.probe = EncoderProbe(model_id) self.probe.set_seed() self.probe.scaled_heads.clear(); self.probe._install_iv_hooks() self.probe.load_input(text, pair or None) self.scan_data = None; self.trace_data = None; self.transcript = []; self.cmp = None self.loaded_model, self.loaded_text, self.loaded_pair = model_id, text, pair return True def ensure(self, model_id, text=""): if self.probe is None: self.load(model_id, text, None) def panes(s): return (render_summary(s), render_input(s.probe), render_headmap(s.probe, s.scan_data), render_prediction(s.probe), render_trace(s.probe, s.trace_data), status_html(s)) def _refresh(s, rescan=True): if s.probe is None or s.probe.input_ids is None: return s.probe.forward_snapshot() if rescan and s.scan_data is not None: layers = sorted(s.scan_data["rows"].keys()) heads = s.scan_data["heads"] or None s.scan_data = s.probe.scan_heads(target_layers=layers, target_heads=heads) # ── command handling ──────────────────────────────────────────────────────── def initial_load(model_id, text, pair, s): rebuilt = s.load(model_id, text, pair or None) _refresh(s, rescan=False) s.scan_data = s.probe.scan_heads() tier = "Tier-2 (per-head)" if s.probe.head_scan_supported else "Tier-1 only" s.transcript.append(("(load)", f"{'loaded' if rebuilt else 'ready'} {s.probe.arch_name} " f"{s.probe.num_layers}L×{s.probe.num_heads}H · " f"labels {','.join(s.probe.label_codes)} · {tier}")) return panes(s) def _resolve_class_arg(p, arg): arg = arg.strip() if arg.isdigit(): return int(arg) al = arg.lower() for i, (n, c) in enumerate(zip(p.label_names, p.label_codes)): if al == c.lower() or n.lower().startswith(al) or al in n.lower(): return i return None def handle(cmd, s, model_id): cmd = (cmd or "").strip() if not cmd: return panes(s) s.ensure(model_id) p = s.probe st = "ok" try: lc = cmd.lower() if lc.startswith("h "): parts = cmd.split() tl = None if parts[1] == "*" else [int(parts[1])] th = None if len(parts) > 2: th = None if parts[2] == "*" else [int(parts[2])] if p.snapshot is None: p.forward_snapshot() s.scan_data = p.scan_heads(target_layers=tl, target_heads=th) st = f"scan layers={parts[1]} heads={parts[2] if len(parts)>2 else '*'}" elif lc.startswith("pos "): p.pos = max(0, int(cmd[4:].strip())) _refresh(s); st = f"pos={p.pos}" elif lc.startswith("spark "): rest = cmd[6:].strip().split() tok = rest[0]; rl, rh = 0, None if len(rest) > 1 and ":" in rest[1]: a, b = rest[1].split(":"); rl, rh = int(a), int(b) if p.snapshot is None: p.forward_snapshot() td = p.vocab_rank_trace(tok, rl, rh) if td == "unsupported": st = f"rank trace unsupported ({p.arch_name}, Tier-1)" elif td is None: st = f"token '{tok}' not in vocab" else: s.trace_data = td; st = f"spark '{tok}'" elif lc.startswith("class "): ci = _resolve_class_arg(p, cmd[6:]) if ci is None: st = f"unknown class '{cmd[6:].strip()}' (have {','.join(p.label_codes)})" else: if p.snapshot is None: p.forward_snapshot() td = p.class_prob_trace(ci) if td == "unsupported": st = f"class trace unsupported ({p.arch_name}, Tier-1)" else: s.trace_data = td; st = f"class {p.label_names[ci]}" elif lc.startswith("mute "): a = cmd.split(); st = p.scale_head(int(a[1]), int(a[2]), 0.0); _refresh(s) elif lc.startswith("scale "): a = cmd.split(); st = p.scale_head(int(a[1]), int(a[2]), float(a[3])); _refresh(s) elif lc.startswith("unmute "): a = cmd.split(); st = p.scale_head(int(a[1]), int(a[2]), 1.0); _refresh(s) elif lc in ("clear", "unmute all"): st = p.clear_interventions(); _refresh(s) elif lc.startswith("cmp "): other = cmd[4:].strip() base_ids, base_mask, base_pos = p.input_ids, p.attn_mask, p.pos base_pred = p.snapshot.pred if p.snapshot else None base_probs = p.snapshot.class_probs.clone() if p.snapshot else None p.load_input(other); p.forward_snapshot() diff = [] if base_probs is not None: for i in range(p.num_labels): d = float(p.snapshot.class_probs[i] - base_probs[i]) diff.append(f"{p.label_codes[i]}{d*100:+.0f}") new_pred = p.label_names[p.snapshot.pred] old_pred = p.label_names[base_pred] if base_pred is not None else "?" s.cmp = (new_pred, diff) # restore the primary trace p.input_ids, p.attn_mask, p.pos = base_ids, base_mask, base_pos p.forward_snapshot() if s.scan_data is not None: s.scan_data = p.scan_heads(target_layers=sorted(s.scan_data["rows"].keys()), target_heads=s.scan_data["heads"] or None) st = f"cmp: {old_pred}→{new_pred} Δ[{' '.join(diff)}]" elif lc == "s": from datetime import datetime fn = f"/tmp/bertographer_{datetime.now():%Y%m%d_%H%M%S}.json" try: with open(fn, "w") as f: json.dump(p.full_log, f, indent=2) st = f"saved {fn}" except Exception as e: st = f"save failed: {e}" elif lc in ("r", "refresh"): _refresh(s); st = "refreshed" elif lc in ("help", "?"): st = "see help line under the command bar" else: st = f"unknown: {cmd}" except Exception as e: st = f"error: {type(e).__name__}: {e}" s.transcript.append((cmd, st)) return panes(s) # ── Gradio UI ─────────────────────────────────────────────────────────────── CSS = """ .gradio-container { max-width: 100% !important; padding: 6px 10px !important; } .gradio-container .main { padding: 0 !important; gap: 4px !important; } .gradio-container .gap, .gradio-container .form { gap: 4px !important; } .gradio-container .block { padding: 3px !important; border-radius: 4px !important; } .gradio-container .prose { margin: 0 !important; } #hdr h1 { font-size: 18px !important; margin: 0 !important; line-height: 1.2 !important; } #hdr p, #hdr em { font-size: 11px !important; margin: 0 !important; color: #8b949e !important; line-height: 1.25 !important; } #hdr { margin: 0 !important; padding: 4px 0 !important; } #topbar { gap: 6px !important; } #topbar textarea, #topbar input[type="text"] { font-size: 12px !important; padding: 4px 6px !important; min-height: 28px !important; line-height: 1.3 !important; } #topbar label, #topbar span[data-testid="block-label"] { font-size: 10px !important; padding: 0 0 2px 0 !important; margin: 0 !important; text-transform: uppercase; letter-spacing: .5px; color: #8b949e !important; } #topbar button { min-height: 56px !important; align-self: stretch !important; font-size: 13px !important; } #cmdbar textarea, #cmdbar input[type="text"] { font-size: 12px !important; padding: 4px 8px !important; min-height: 28px !important; font-family: 'JetBrains Mono','Fira Code',Consolas,monospace !important; } #cmdbar label, #cmdbar span[data-testid="block-label"] { font-size: 10px !important; color: #8b949e !important; margin: 0 !important; } #cmdbar button { min-height: 36px !important; font-size: 12px !important; } #help { font-size: 10.5px !important; color: #8b949e !important; margin: 2px 0 !important; } #help code { font-size: 10.5px !important; padding: 0 2px !important; background: #161b22 !important; } .pane { border: 1px solid #30363d; border-radius: 4px; background: #0d1117; color: #c9d1d9; font-family: 'JetBrains Mono','Fira Code',Consolas,monospace; font-size: 12px; } .pane h3 { margin: 0; padding: 3px 8px; background: #161b22; border-bottom: 1px solid #30363d; font-size: 10px; letter-spacing: .5px; text-transform: uppercase; color: #8b949e; } .pane-body { padding: 5px 8px; max-height: 460px; overflow: auto; } .dim { color: #6e7681; } .context-text { white-space: pre-wrap; word-break: break-word; margin-top: 4px; line-height: 1.7; } .tokn { background: #11161c; border: 1px solid #21262d; border-radius: 3px; padding: 0 3px; } .pos-here { background: #0f2418; outline: 1px solid #64D32A; color: #b8f0c4; border-radius: 3px; padding: 0 3px; font-weight: bold; } table.headmap { border-collapse: collapse; font-size: 11px; } table.headmap th, table.headmap td { padding: 1px 4px; border: 1px solid #21262d; text-align: center; white-space: pre; } table.headmap th.ps { background: #112; } table.headmap th.full { background: #133; color: #64D32A; } table.headmap td.lay { color: #8b949e; } table.headmap td.full { background: #0a1410; } table.headmap tr.subhdr th { background: #161b22; color: #8b949e; font-weight: normal; font-size: 10px; } .hm-cell .cls { font-weight: bold; padding-left: 3px; } .hm-cell .cos { color: #6e7681; padding-left: 2px; } .legend { padding: 4px 0 0 0; font-size: 10px; } .predline { padding: 2px 0 6px 0; font-size: 13px; } table.probs { border-collapse: collapse; font-size: 12px; width: 100%; } table.probs td { padding: 1px 6px; white-space: pre; } table.probs td.plabel { font-weight: bold; } table.probs td.pbar { font-family: monospace; letter-spacing: -1px; } table.probs td.ppct { text-align: right; color: #c9d1d9; } .votes { margin-top: 8px; font-size: 11px; line-height: 1.6; } .vote { display: inline-block; min-width: 22px; text-align: center; background: #11161c; border: 1px solid #21262d; border-radius: 3px; margin: 1px; padding: 0 2px; font-weight: bold; } #summary-strip { padding: 0 !important; margin: 0 !important; } .summary { font-family: 'JetBrains Mono',Consolas,monospace; font-size: 11px; padding: 4px 8px; background: #0d1117; border: 1px solid #30363d; border-radius: 4px; color: #c9d1d9; line-height: 1.5; } .summary .sep { color: #30363d; padding: 0 4px; } .summary .dim { color: #6e7681; } #dual-row { align-items: stretch !important; } #dual-row > div { display: flex; flex-direction: column; } #dual-row #col-pred, #dual-row #col-head { height: 56vh; min-height: 360px; max-height: 660px; } #dual-row .pane { height: 100%; display: flex; flex-direction: column; overflow: hidden; } #dual-row .pane h3 { flex: 0 0 auto; } #dual-row .pane .pane-body { flex: 1 1 auto; max-height: none !important; overflow: auto; } pre.spark { margin: 0; font-size: 11px; line-height: 1.1; } .prompt { color: #64D32A; } .gradio-container textarea::placeholder, .gradio-container input[type="text"]::placeholder { color: #6e7681 !important; opacity: 1 !important; } """ HELP = ( "`h *` / `h L` / `h L H` head×layer scan · `pos N` inspect token position · " "`spark [lo:hi]` vocab-rank trace · `class