Spaces:
Running on Zero
Running on Zero
| """Running a model and reading the cone off its own trajectory. | |
| rt = ulens.Runner("qwen3.5-4b") | |
| tr = rt.run("Is Hamlet's madness feigned?", max_new_tokens=600, seed=42) | |
| tr.u_lens # the score | |
| tr.field # [n_blocks, T, k] signed projections: depth x token x axis | |
| tr.a_cone_series() # A_cone per token at the readout block | |
| tr.surface() # A_cone over depth x token | |
| Two passes: `generate` with the model's own sampling settings, then one teacher-forced pass | |
| over prompt + generation that reads every lens block through a forward hook and the per-token | |
| entropies through the unembedding. The second pass is what makes the field exact — a | |
| generation pass sees each token once, with the state it had at the time, which is the same | |
| thing only when nothing is being intervened on. | |
| Alignment, which decides whether the score matches the published one. `entropy[i]` is the | |
| entropy of the distribution that **produced** token `i`, read from the logits at `i - 1`, so | |
| the thinking-block mean is `entropy[prompt_len : eot + 1].mean()` and includes the | |
| distribution that produced the end-of-think token itself. `A_cone` is read at the end-of-think | |
| token's own state. Both match `pipeline/04_forward.py` of the reproduction code. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import re | |
| from dataclasses import dataclass | |
| from dataclasses import field as _field | |
| from pathlib import Path | |
| import numpy as np | |
| from . import axes as _axes | |
| from . import basis as _basis | |
| from . import readout as _readout | |
| from ._deps import require, require_extra | |
| from .models import ModelSpec | |
| from .models import get as _model | |
| SEG_PROMPT, SEG_THINK, SEG_ANSWER = "prompt", "think", "answer" | |
| #: positions per chunk when pushing hidden states through the unembedding | |
| LOGIT_CHUNK = 256 | |
| class Trace: | |
| """One generation, with the cone read off every depth and every token.""" | |
| model: str | |
| prompt: str | |
| tokens: list[int] | |
| token_strs: list[str] | |
| prompt_len: int | |
| eot: int # absolute index of the end-of-think token | |
| entropy: np.ndarray # [T], of the distribution that produced each token | |
| field: np.ndarray # [n_blocks, T, k] signed projections | |
| blocks: list[int] | |
| names: list[str] | |
| readout_block: int | |
| generated: str | |
| truncated: bool # the thinking block was never closed | |
| config: dict = _field(default_factory=dict) | |
| # ------------------------------------------------------------------ geometry | |
| def n_tokens(self) -> int: | |
| return len(self.tokens) | |
| def n_axes(self) -> int: | |
| return int(self.field.shape[2]) | |
| def think_len(self) -> int: | |
| return self.eot + 1 - self.prompt_len | |
| def segments(self) -> np.ndarray: | |
| """`[T]` of 'prompt' / 'think' / 'answer'.""" | |
| s = np.array([SEG_ANSWER] * self.n_tokens, dtype=object) | |
| s[: self.prompt_len] = SEG_PROMPT | |
| s[self.prompt_len: self.eot + 1] = SEG_THINK | |
| return s | |
| def block_index(self, block: int | None = None) -> int: | |
| b = self.readout_block if block is None else block | |
| if b not in self.blocks: | |
| raise KeyError(f"block {b} was not captured; have {self.blocks}") | |
| return self.blocks.index(b) | |
| def projections(self, block: int | None = None) -> np.ndarray: | |
| """`[T, k]` signed projections at one block.""" | |
| return self.field[self.block_index(block)] | |
| def components(self, block: int | None = None) -> np.ndarray: | |
| """`[T, k]` rectified components at one block.""" | |
| return _readout.rectify(self.projections(block)) | |
| def a_cone_series(self, block: int | None = None) -> np.ndarray: | |
| """`[T]` A_cone per token at one block.""" | |
| return _readout.a_cone_of(self.projections(block)) | |
| def surface(self) -> np.ndarray: | |
| """`[n_blocks, T]` A_cone over depth and token: the field, as one number per cell.""" | |
| return _readout.a_cone_of(self.field) | |
| # ------------------------------------------------------------------ the score | |
| def a_cone(self) -> float: | |
| """A_cone at the end-of-think token, at the readout block.""" | |
| return float(self.a_cone_series()[self.eot]) | |
| def think_entropy(self) -> float: | |
| """Mean entropy of the distributions that produced the thinking block's tokens; | |
| NaN when there is no thinking block (thinking off), without numpy's warning.""" | |
| e = self.entropy[self.prompt_len: self.eot + 1] | |
| return float(e.mean()) if e.size else float("nan") | |
| def u_lens(self) -> float: | |
| """The score. Higher means more likely to be wrong. | |
| Check `readout_valid` before comparing this with a published number. | |
| """ | |
| return self.a_cone * self.think_entropy | |
| def readout_valid(self) -> bool: | |
| """Whether the score was read where the method says to read it. | |
| False when the thinking block was never closed: the readout then falls on the last | |
| token of a budget-truncated block rather than on an end-of-think marker, and the | |
| answer is empty. The number is still computed — the field is real — but it is not | |
| the published quantity and must not be compared with one. | |
| """ | |
| return not self.truncated and self.think_len > 0 | |
| def summary(self) -> str: | |
| t = ("TRUNCATED, never closed - readout is not the published one" | |
| if self.truncated else "closed") | |
| return (f"{self.model}: {self.n_tokens} tokens " | |
| f"(prompt {self.prompt_len}, think {self.think_len} [{t}], " | |
| f"answer {self.n_tokens - self.eot - 1}) · " | |
| f"{len(self.blocks)} blocks [{self.blocks[0]}-{self.blocks[-1]}] · " | |
| f"readout block {self.readout_block}\n" | |
| f" A_cone {self.a_cone:.4f} × think entropy {self.think_entropy:.4f} " | |
| f"= U-Lens {self.u_lens:.4f}") | |
| def axis_table(self, block: int | None = None) -> str: | |
| """Per-axis mean rectified component over the thinking block.""" | |
| C = self.components(block)[self.prompt_len: self.eot + 1] | |
| tot = float((C ** 2).sum()) or 1.0 | |
| b = self.readout_block if block is None else block | |
| out = [f"axis contributions over the thinking block (block {b})"] | |
| for i, n in enumerate(self.names): | |
| out.append(f" {n:22s} mean {C[:, i].mean():.4f} " | |
| f"share {100 * (C[:, i] ** 2).sum() / tot:5.1f}% " | |
| f"argmax {100 * (C.argmax(1) == i).mean():5.1f}%") | |
| return "\n".join(out) | |
| # ------------------------------------------------------------------ persistence | |
| def _meta(self) -> dict: | |
| return {"model": self.model, "prompt": self.prompt, "prompt_len": self.prompt_len, | |
| "eot": self.eot, "blocks": list(self.blocks), "names": list(self.names), | |
| "readout_block": self.readout_block, "generated": self.generated, | |
| "truncated": self.truncated, "token_strs": self.token_strs, | |
| "config": self.config} | |
| def save(self, path: str | Path) -> Path: | |
| """Write `<path>.npz` plus a readable `<path>.json` sidecar.""" | |
| path = Path(path) | |
| if path.suffix != ".npz": | |
| path = path.with_suffix(".npz") | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| np.savez_compressed(path, field=self.field.astype(np.float32), | |
| entropy=self.entropy.astype(np.float32), | |
| tokens=np.array(self.tokens), | |
| meta=np.array(json.dumps(self._meta()))) | |
| side = dict(self._meta()) | |
| # a trace without a thinking block has no thinking-block entropy and no score; the | |
| # sidecar says null, which is JSON, rather than NaN, which is not | |
| side["scores"] = {k: (float(v) if np.isfinite(v) else None) for k, v in | |
| (("a_cone", self.a_cone), ("think_entropy", self.think_entropy), | |
| ("u_lens", self.u_lens))} | |
| path.with_suffix(".json").write_text(json.dumps(side, indent=1, allow_nan=False) | |
| + "\n") | |
| return path | |
| def load(path: str | Path) -> Trace: | |
| """Read a trace back. No model and no GPU needed — every view works from the field.""" | |
| z = np.load(Path(path), allow_pickle=False) | |
| m = json.loads(str(z["meta"])) | |
| return Trace(model=m["model"], prompt=m["prompt"], tokens=[int(t) for t in z["tokens"]], | |
| token_strs=m["token_strs"], prompt_len=m["prompt_len"], eot=m["eot"], | |
| entropy=np.asarray(z["entropy"], np.float32), | |
| field=np.asarray(z["field"], np.float32), blocks=m["blocks"], | |
| names=m["names"], readout_block=m["readout_block"], | |
| generated=m["generated"], truncated=m["truncated"], | |
| config=m.get("config", {})) | |
| class Runner: | |
| """A loaded model, its cone at every lens depth, and the pass that reads them together.""" | |
| def __init__(self, model: str | ModelSpec, device: str = "cuda", dtype=None, | |
| blocks: list[int] | None = None, attn_implementation: str = "sdpa", | |
| device_map=None, verbose: bool = True): | |
| require_extra("run") | |
| torch = require("torch") | |
| self.spec = _model(model) | |
| self.device = device | |
| self.verbose = verbose | |
| self._log(f"building cones at every lens block for {self.spec.key} …") | |
| self.stack: _axes.ConeStack = _basis.field(self.spec, blocks=blocks, | |
| verbose=verbose) | |
| self._log(f"loading {self.spec.hf} …") | |
| AutoModelForCausalLM = require("transformers").AutoModelForCausalLM | |
| AutoTokenizer = require("transformers").AutoTokenizer | |
| kw = {"revision": self.spec.revision} if self.spec.revision else {} | |
| self.tok = AutoTokenizer.from_pretrained(self.spec.hf, **kw) | |
| self.model = AutoModelForCausalLM.from_pretrained( | |
| self.spec.hf, dtype=dtype or torch.bfloat16, | |
| device_map=device_map if device_map is not None else device, | |
| attn_implementation=attn_implementation, **kw).eval() | |
| self._layers = _decoder_blocks(self.model) | |
| self.open_id, self.close_id = _think_ids(self.tok, self.spec) | |
| self._U = {b: torch.tensor(self.stack.U[i], device=self._param_device()) | |
| for i, b in enumerate(self.stack.blocks)} | |
| self._log(f"ready: {len(self.stack.blocks)} blocks, {self.stack.n_axes} axes " | |
| f"({', '.join(self.stack.names)}), think ids " | |
| f"open={self.open_id} close={self.close_id}") | |
| def _log(self, msg: str) -> None: | |
| if self.verbose: | |
| print(msg, flush=True) | |
| def _param_device(self): | |
| return next(self.model.parameters()).device | |
| # ------------------------------------------------------------------ prompting | |
| def prompt_ids(self, prompt: str, thinking: bool = True) -> list[int]: | |
| """Prompt token ids through the model's own chat template.""" | |
| msgs = [{"role": "user", "content": prompt}] | |
| for kw in ({"enable_thinking": bool(thinking)}, {}): | |
| try: | |
| ids = self.tok.apply_chat_template(msgs, tokenize=True, | |
| add_generation_prompt=True, | |
| return_dict=False, **kw) | |
| except TypeError: | |
| continue | |
| return _flatten(ids) | |
| raise RuntimeError("the chat template rejected both call forms") | |
| # ------------------------------------------------------------------ the passes | |
| def generate(self, ids: list[int], max_new_tokens: int = 1024, seed: int = 0, | |
| **sampling) -> list[int]: | |
| """Sample a continuation. Returns prompt + generation.""" | |
| torch = require("torch") | |
| cfg = dict(self.spec.sampling) | |
| cfg.update(sampling) | |
| x = torch.tensor([ids], device=self._param_device()) | |
| torch.manual_seed(int(seed)) | |
| with torch.no_grad(): | |
| out = self.model.generate(x, attention_mask=torch.ones_like(x), | |
| max_new_tokens=int(max_new_tokens), do_sample=True, | |
| pad_token_id=self.tok.eos_token_id, **cfg) | |
| return [int(t) for t in out[0]] | |
| def capture(self, ids: list[int], prompt_len: int, prompt: str = "", | |
| config: dict | None = None) -> Trace: | |
| """One teacher-forced pass: the field at every lens block, and every entropy.""" | |
| torch = require("torch") | |
| dev = self._param_device() | |
| proj: dict[int, np.ndarray] = {} | |
| hooks = [] | |
| def mk(b): | |
| def hook(_m, _i, o): | |
| h = (o[0] if isinstance(o, tuple) else o)[0].float() | |
| h = h / (h.norm(dim=-1, keepdim=True) + 1e-9) | |
| proj[b] = (h @ self._U[b]).cpu().numpy().astype(np.float32) | |
| return hook | |
| for b in self.stack.blocks: | |
| hooks.append(self._layers[b].register_forward_hook(mk(b))) | |
| x = torch.tensor([ids], device=dev) | |
| try: | |
| with torch.no_grad(): | |
| h = _decoder(self.model)(input_ids=x, use_cache=False).last_hidden_state[0] | |
| ent = self._entropies(h) | |
| finally: | |
| for hk in hooks: | |
| hk.remove() | |
| del h | |
| torch.cuda.empty_cache() if dev.type == "cuda" else None | |
| F = np.stack([proj[b] for b in self.stack.blocks], 0) # [L, T, k] | |
| eot, truncated = _segment(ids, prompt_len, self.open_id, self.close_id) | |
| strs = [self.tok.decode([t]) for t in ids] | |
| return Trace(model=self.spec.key, prompt=prompt, tokens=list(ids), token_strs=strs, | |
| prompt_len=prompt_len, eot=eot, entropy=ent, field=F, | |
| blocks=list(self.stack.blocks), names=list(self.stack.names), | |
| readout_block=self.spec.readout_block, | |
| generated=self.tok.decode(ids[prompt_len:]), truncated=truncated, | |
| config=dict(config or {})) | |
| def _entropies(self, h) -> np.ndarray: | |
| """`[T]` entropy of the distribution that produced each token; position 0 is NaN. | |
| The distribution for token `i` lives at position `i - 1`, so the array is shifted by | |
| one and the first token, which nothing predicted, has no entropy. | |
| """ | |
| torch = require("torch") | |
| T = int(h.shape[0]) | |
| out = np.full(T, np.nan, np.float32) | |
| head = self.model.get_output_embeddings() | |
| cap = getattr(_text_config(self.model), "final_logit_softcapping", None) | |
| with torch.no_grad(): | |
| for i in range(0, T - 1, LOGIT_CHUNK): | |
| j = min(i + LOGIT_CHUNK, T - 1) | |
| z = head(h[i:j].to(self.model.dtype)).float() | |
| if cap: | |
| z = torch.tanh(z / cap) * cap | |
| g = torch.log_softmax(z, dim=-1) | |
| out[i + 1:j + 1] = (-(g.exp() * g).sum(-1)).cpu().numpy() | |
| return out | |
| def run(self, prompt: str, max_new_tokens: int = 1024, seed: int = 0, | |
| thinking: bool = True, **sampling) -> Trace: | |
| """Generate from `prompt`, then capture the field over the whole sequence.""" | |
| pre = self.prompt_ids(prompt, thinking) | |
| self._log(f"generating up to {max_new_tokens} tokens from {len(pre)} prompt tokens …") | |
| ids = self.generate(pre, max_new_tokens=max_new_tokens, seed=seed, **sampling) | |
| self._log(f"teacher-forced pass over {len(ids)} tokens × " | |
| f"{len(self.stack.blocks)} blocks …") | |
| cfg = dict(self.spec.sampling) | |
| cfg.update(sampling) | |
| return self.capture(ids, len(pre), prompt=prompt, | |
| config={"seed": seed, "max_new_tokens": max_new_tokens, | |
| "thinking": thinking, "sampling": cfg, | |
| "hf": self.spec.hf, "revision": self.spec.revision}) | |
| # ---------------------------------------------------------------------- helpers | |
| def _flatten(ids) -> list[int]: | |
| if hasattr(ids, "input_ids"): | |
| ids = ids.input_ids | |
| ids = list(ids) | |
| if ids and isinstance(ids[0], (list, tuple)): | |
| ids = list(ids[0]) | |
| return [int(i) for i in ids] | |
| def _text_config(model): | |
| return getattr(model.config, "text_config", model.config) | |
| def _decoder(model): | |
| return getattr(model, "model", model) | |
| def _decoder_blocks(model) -> dict[int, object]: | |
| """`{index: block}` for the deepest `.layers.N` group, which is the decoder stack.""" | |
| groups: dict[str, dict[int, object]] = {} | |
| for name, mod in model.named_modules(): | |
| m = re.match(r"(.*)\.layers\.(\d+)$", name) | |
| if m: | |
| groups.setdefault(m.group(1), {})[int(m.group(2))] = mod | |
| if not groups: | |
| raise RuntimeError("no `.layers.N` modules found; cannot locate the decoder blocks") | |
| return groups[max(groups, key=lambda k: len(groups[k]))] | |
| def _think_ids(tok, spec: ModelSpec) -> tuple[int | None, int | None]: | |
| """(open, close) token ids of the thinking channel, or None when not single tokens.""" | |
| def one(s: str) -> int | None: | |
| ids = tok.encode(s, add_special_tokens=False) | |
| return int(ids[0]) if len(ids) == 1 else None | |
| return one(spec.think_start), one(spec.think_end) | |
| def _segment(ids: list[int], prompt_len: int, open_id: int | None, close_id: int | None, | |
| ) -> tuple[int, bool]: | |
| """(absolute index of the end-of-think token, whether the block was never closed). | |
| The chat template may already have opened the thinking channel at the end of the prompt, | |
| which is the usual case; otherwise the opener is looked for in the generation. A block | |
| that is opened and never closed runs to the last token, and the whole generation is | |
| thinking — the case a token budget produces, and the one that must not be read as a | |
| zero-token answer. | |
| """ | |
| head, gen = ids[:prompt_len], ids[prompt_len:] | |
| if close_id is None: | |
| return len(ids) - 1, True | |
| last_open = max((i for i, t in enumerate(head) if t == open_id), default=-1) | |
| last_close = max((i for i, t in enumerate(head) if t == close_id), default=-1) | |
| if open_id is not None and last_open > last_close: | |
| start = 0 # already open at the prompt's end | |
| else: | |
| start = next((i + 1 for i, t in enumerate(gen) if t == open_id), None) | |
| if start is None: | |
| return max(prompt_len - 1, 0), False # no block at all: it is all answer | |
| rest = gen[start:] | |
| if close_id in rest: | |
| return prompt_len + start + rest.index(close_id), False | |
| return len(ids) - 1, True | |