Spaces:
Running
Running
fix: cpu-basic + gradio5 pin + module-level demo + ungated default + no import-time cuda
afbece7 verified | """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()) | |
| 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 = [] | |
| 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 "<div class='pane-body'><i>No trace loaded. Enter text and click Load.</i></div>" | |
| toks = p.tokenizer.convert_ids_to_tokens(p.input_ids[0]) | |
| special = set(p.tokenizer.all_special_tokens) | |
| parts = [f"<div class='dim'>[{len(toks)} tokens · inspecting pos {p.pos}]</div>", | |
| "<div class='context-text'>"] | |
| for i, t in enumerate(toks): | |
| disp = html_escape(t.replace("▁", "·").replace("Ġ", "·")) | |
| if i == p.pos: | |
| parts.append(f"<span class='pos-here'>[{i}:{disp}]</span> ") | |
| elif t in special: | |
| parts.append(f"<span class='dim'>[{i}:{disp}]</span> ") | |
| else: | |
| parts.append(f"<span class='tokn'><span class='dim'>{i}:</span>{disp}</span> ") | |
| parts.append("</div>") | |
| return "<div class='pane-body scroll'>" + "".join(parts) + "</div>" | |
| 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"<td class='hm-cell'>" | |
| f"<span style='color:{tc}'>{html_escape(fmt_tok(tok, CELL_W))}</span>" | |
| f"<span class='cls' style='color:{cc}'>{code}{int(cp*100):02d}</span>" | |
| f"<span class='cos'>.{int(abs(cell['cosine'])*99):02d}</span></td>") | |
| def render_headmap(p: EncoderProbe, data: dict) -> str: | |
| if not data: | |
| return ("<div class='pane-body'><i>No scan yet. <code>h *</code> scans all " | |
| "layers, <code>h L</code> one layer.</i></div>") | |
| rows_data = data["rows"] | |
| heads = data["heads"] | |
| layers = sorted(rows_data.keys()) | |
| if not layers: | |
| return "<div class='pane-body'><i>Empty scan.</i></div>" | |
| hdr = ["<tr><th class='lay'>Lay</th>"] | |
| for h in heads: hdr.append(f"<th class='ps'>H{h}</th>") | |
| hdr.append("<th class='full'>Layer</th></tr>") | |
| sub = ["<tr class='subhdr'><th></th>"] | |
| if heads: | |
| sub.append(f"<th class='ps' colspan='{len(heads)}'>per-head span " | |
| f"(head→o_proj→vocab|class)</th>") | |
| sub.append("<th class='full'>residual</th></tr>") | |
| body = [] | |
| for l in layers: | |
| r = [f"<tr><td class='lay'>L{l}</td>"] | |
| row = rows_data[l] | |
| for h in heads: | |
| r.append(_cell(p, row[h]) if h in row else "<td class='hm-cell dim'>·</td>") | |
| r.append(_cell(p, row["Layer"]).replace("hm-cell", "hm-cell full") | |
| if "Layer" in row else "<td class='hm-cell full dim'>·</td>") | |
| r.append("</tr>") | |
| body.append("".join(r)) | |
| note = "" | |
| if not heads: | |
| note = ("<div class='dim' style='padding:2px 0'>per-head decomposition unavailable " | |
| "for this architecture — full-layer residual only (Tier-1)</div>") | |
| return (f"<div class='pane-body scroll'>{note}" | |
| f"<table class='headmap'>{''.join(sub)}{''.join(hdr)}{''.join(body)}</table>" | |
| f"<div class='legend dim'>cell: <b>tok</b> (logit-lens) · " | |
| f"<b>CODE pp</b> (head's class vote) · <b>.cos</b> (cosine to that class)</div></div>") | |
| def render_prediction(p: EncoderProbe) -> str: | |
| if p is None or p.snapshot is None: | |
| return "<div class='pane-body'><i>Load a trace to see the prediction.</i></div>" | |
| s = p.snapshot | |
| pred = s.pred | |
| rows = ["<div class='pane-body scroll'>"] | |
| rows.append(f"<div class='predline'>prediction " | |
| f"<b style='color:{class_color(p.label_names[pred], pred)}'>" | |
| f"{html_escape(p.label_names[pred])}</b> " | |
| f"<span class='dim'>({s.class_probs[pred]*100:.1f}%) · " | |
| f"margin {s.margin*100:.1f}pp · H={s.entropy:.2f} bits</span></div>") | |
| rows.append("<table class='probs'>") | |
| 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"<tr><td class='plabel' style='color:{c}'>{html_escape(p.label_names[i])}</td>" | |
| f"<td class='pbar'><span style='color:{c}'>{'█'*bar}{'·'*(28-bar)}</span></td>" | |
| f"<td class='ppct'>{pr*100:5.1f}%{mark}</td></tr>") | |
| rows.append("</table>") | |
| # 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"<span class='vote' style='color:{c}' title='H{hi}'>" | |
| f"{p.label_codes[ci]}</span>") | |
| rows.append(f"<div class='votes'><span class='dim'>final-layer head votes " | |
| f"(L{li}):</span> {' '.join(cells)}</div>") | |
| rows.append("</div>") | |
| return "".join(rows) | |
| def render_trace(p: EncoderProbe, td: dict) -> str: | |
| if not td: | |
| return ("<div class='pane-body'><i>Trace one signal across heads×layers: " | |
| "<code>spark <token></code> (vocab rank) or " | |
| "<code>class <label></code> (vote probability).</i></div>") | |
| 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"<span class='dim'>{html_escape(hdr)}</span>") | |
| 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"<span class='dim'>{html_escape(hdr)}</span>") | |
| 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"<div class='pane-body scroll'><pre class='spark'>{chr(10).join(lines)}</pre></div>" | |
| 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"<span style='color:{color}'> {ch} </span>") | |
| else: | |
| cells.append(f" {ch} ") | |
| lines.append(html_escape(lab) + "".join(cells)) | |
| def render_summary(s) -> str: | |
| if s.probe is None: | |
| return "<div class='summary dim'>no model loaded</div>" | |
| 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"<b style='color:{class_color(p.label_names[snap.pred], snap.pred)}'>" | |
| f"{html_escape(p.label_names[snap.pred])}</b> {snap.class_probs[snap.pred]*100:.0f}%" | |
| if snap else "—") | |
| iv = (f"<span style='color:#ff6b6b'>iv {len(p.scaled_heads)}</span>" | |
| if p.scaled_heads else "iv 0") | |
| items = [ | |
| f"<b>{html_escape(p.arch_name)}</b> <span class='dim'>{html_escape(p.model_id)}</span>", | |
| f"<span class='dim'>tier</span> {tier}", | |
| f"<span class='dim'>L</span>{p.num_layers} <span class='dim'>×H</span>{p.num_heads}", | |
| f"<span class='dim'>tok</span> {n_tok}", | |
| f"<span class='dim'>pos</span> {p.pos}", | |
| f"<span class='dim'>pred</span> {pred}", | |
| f"<span class='dim'>H</span>=<b>{snap.entropy:.2f}</b>" if snap else "", | |
| iv, | |
| ] | |
| return "<div class='summary'>" + " <span class='sep'>·</span> ".join(x for x in items if x) + "</div>" | |
| def status_html(s) -> str: | |
| if not s.transcript: | |
| return "<div class='pane-body dim'>Ready.</div>" | |
| rows = [] | |
| for cmd, st in s.transcript[-22:]: | |
| rows.append(f"<div><span class='prompt'>›</span> <code>{html_escape(cmd)}</code> " | |
| f"<span class='dim'>— {html_escape(st)}</span></div>") | |
| return f"<div class='pane-body scroll'>{''.join(rows)}</div>" | |
| # ── 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 <tok> [lo:hi]` vocab-rank trace · `class <label>` per-head vote trace · " | |
| "`mute L H` · `scale L H X` · `unmute L H` · `clear` · " | |
| "`cmp <text>` diff vs another trace · `r` refresh · `s` save" | |
| ) | |
| MODEL_PRESETS = [ | |
| "cardiffnlp/twitter-roberta-base-sentiment-latest", | |
| "cardiffnlp/twitter-roberta-base-emotion", | |
| "distilbert-base-uncased-finetuned-sst-2-english", | |
| "textattack/bert-base-uncased-SST-2", | |
| "MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli", | |
| "google/electra-base-discriminator", | |
| "ProsusAI/finbert", | |
| ] | |
| def build_ui(): | |
| with gr.Blocks(title="Bertographer", css=CSS, theme=gr.themes.Base()) as demo: | |
| gr.Markdown( | |
| "# Bertographer \n*Mechanistic cockpit for encoder **classifiers** — " | |
| "per-head logit-lens · per-head class votes · vocab/class rank traces · " | |
| "head interventions — on BERT / RoBERTa / DeBERTa / DistilBERT / ELECTRA. " | |
| "An input is a **trace**; layers are the **waterfall**; heads are **voters**; " | |
| "muting a head is the **counterfactual**.*", | |
| elem_id="hdr") | |
| session = gr.State(Session()) | |
| with gr.Row(elem_id="topbar"): | |
| model_id = gr.Dropdown(choices=MODEL_PRESETS, value=DEFAULT_MODEL_ID, label="Model", | |
| allow_custom_value=True, filterable=True, scale=4) | |
| text = gr.Textbox(label="Text (trace)", | |
| value="I can't believe how good this turned out — absolutely thrilled.", | |
| lines=1, max_lines=3, scale=6) | |
| pair = gr.Textbox(label="Text pair (NLI/optional)", value="", lines=1, max_lines=3, scale=4) | |
| load_btn = gr.Button("Load", variant="primary", scale=1, min_width=80) | |
| summary_pane = gr.HTML(value="<div class='summary dim'>no model loaded</div>", elem_id="summary-strip") | |
| ctx_pane = gr.HTML(value="<div class='pane'><h3>Trace · tokenized input</h3>" | |
| "<div class='pane-body'><i>Click <b>Load</b> to begin.</i></div></div>") | |
| with gr.Row(elem_id="cmdbar"): | |
| cmd = gr.Textbox(label="CMD", placeholder=HELP, scale=10, autofocus=True, | |
| lines=1, max_lines=1) | |
| run_btn = gr.Button("Run", scale=1, variant="primary", min_width=70) | |
| gr.Markdown(HELP, elem_id="help") | |
| with gr.Row(elem_id="dual-row"): | |
| with gr.Column(scale=3, min_width=300, elem_id="col-pred"): | |
| pred_pane = gr.HTML(value="<div class='pane'><h3>Prediction · class probs · head votes</h3>" | |
| "<div class='pane-body'><i>Appears after Load.</i></div></div>") | |
| with gr.Column(scale=7, elem_id="col-head"): | |
| hm_pane = gr.HTML(value="<div class='pane'><h3>Head Map · per-head vocab|class at pos</h3>" | |
| "<div class='pane-body'><i>Auto-scans after each step.</i></div></div>") | |
| trace_pane = gr.HTML(value="<div class='pane'><h3>Rank / Vote Trace · per (head,layer)</h3>" | |
| "<div class='pane-body'><i>Run <code>spark <token></code> or " | |
| "<code>class <label></code>.</i></div></div>") | |
| status = gr.HTML(value="<div class='pane'><h3>Transcript</h3>" | |
| "<div class='pane-body dim'>Ready.</div></div>") | |
| def _wrap(summ, ctx, hm, pred, tr, stt): | |
| return (summ, | |
| f"<div class='pane'><h3>Trace · tokenized input</h3>{ctx}</div>", | |
| f"<div class='pane'><h3>Head Map · per-head vocab|class at pos</h3>{hm}</div>", | |
| f"<div class='pane'><h3>Prediction · class probs · head votes</h3>{pred}</div>", | |
| f"<div class='pane'><h3>Rank / Vote Trace · per (head,layer)</h3>{tr}</div>", | |
| f"<div class='pane'><h3>Transcript</h3>{stt}</div>") | |
| OUT = [summary_pane, ctx_pane, hm_pane, pred_pane, trace_pane, status] | |
| def on_load(mid, t, pr, s): | |
| summ, ctx, hm, pred, tr, stt = initial_load(mid, t, pr, s) | |
| return _wrap(summ, ctx, hm, pred, tr, stt) + (s,) | |
| def on_cmd(c, s, mid): | |
| summ, ctx, hm, pred, tr, stt = handle(c, s, mid) | |
| return _wrap(summ, ctx, hm, pred, tr, stt) + ("", s) | |
| load_btn.click(on_load, inputs=[model_id, text, pair, session], outputs=OUT + [session]) | |
| run_btn.click(on_cmd, inputs=[cmd, session, model_id], outputs=OUT + [cmd, session]) | |
| cmd.submit(on_cmd, inputs=[cmd, session, model_id], outputs=OUT + [cmd, session]) | |
| return demo | |
| # Module-level `demo` so the HF Spaces runner finds it without guessing. | |
| demo = build_ui() | |
| demo.queue() | |
| if __name__ == "__main__": | |
| demo.launch() | |