Spaces:
Running on Zero
Running on Zero
| """Registry of models a U-Space basis can be built for, with their pinned revisions. | |
| The readout block is the model's block at relative depth 2/3. It is recorded explicitly so a | |
| basis is reproducible without fetching a config, and `tests/test_models.py` checks every | |
| entry against `rel_to_block`. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass, field | |
| #: relative depth the state is READ at, one of three pre-registered depths (2/3, 0.75, 0.80) | |
| READOUT_REL = 2 / 3 | |
| #: relative depth the axes DECODE at. Not the same depth: the axes return recognisable words | |
| #: only around rel 0.75-0.80, while the score is best at 2/3. Reading an axis's identity and | |
| #: reading its activation are therefore two different depths, and both are reported as such. | |
| DECODE_REL = 0.75 | |
| #: the published Jacobian-lens collection | |
| LENS_REPO = "neuronpedia/jacobian-lens" | |
| LENS_REV = "0731326edff4ae730ffc5356fe1a4728c748b3a6" | |
| class LensFile: | |
| """Where a model's Jacobian lens file lives. | |
| `repo`/`file` name a file on the Hub; `path` a local `.pt` (all fitted blocks) or `.npy` | |
| (one block, already extracted). | |
| """ | |
| repo: str | None = None | |
| file: str | None = None | |
| revision: str | None = None | |
| path: str | None = None | |
| def __post_init__(self) -> None: | |
| if not self.path and not (self.repo and self.file): | |
| raise ValueError("a LensFile needs either `path` or both `repo` and `file`") | |
| class ModelSpec: | |
| """A model, its pinned revision, its lens, and the block the state is read at.""" | |
| key: str | |
| hf: str | |
| revision: str | |
| n_blocks: int | |
| hidden_size: int | |
| readout_block: int | |
| lens: LensFile | |
| think_end: str | |
| think_start: str = "<think>" | |
| #: the sampling settings the reported runs used | |
| sampling: dict = field(default_factory=lambda: {"temperature": 1.0, "top_p": 0.95, | |
| "top_k": 20}) | |
| aliases: tuple[str, ...] = field(default=()) | |
| paper_model: bool = False | |
| def rel_to_block(rel: float, n_blocks: int) -> int: | |
| """Relative depth -> lens block key. | |
| The lens is keyed by decoder block `b` and transports that block's *output*, i.e. hidden | |
| state index `b + 1`, so `rel = (b + 1) / n_blocks`. | |
| """ | |
| h = round(rel * n_blocks) | |
| return min(max(h, 1), n_blocks - 1) - 1 | |
| MODELS: dict[str, ModelSpec] = {} | |
| def _add(spec: ModelSpec) -> None: | |
| MODELS[spec.key] = spec | |
| _add(ModelSpec( | |
| key="gemma-4-31b-it", | |
| hf="google/gemma-4-31B-it", | |
| revision="842da3794eaa0b77d5f08bae87a17459d91ff475", | |
| n_blocks=60, hidden_size=5376, readout_block=39, | |
| lens=LensFile(repo=LENS_REPO, revision=LENS_REV, | |
| file="gemma-4-31b/jlens/Salesforce-wikitext/gemma-4-31B_jacobian_lens.pt"), | |
| think_end="<channel|>", think_start="<think>", | |
| sampling={"temperature": 1.0, "top_p": 0.95, "top_k": 64}, | |
| aliases=("gemma",), paper_model=True, | |
| )) | |
| _add(ModelSpec( | |
| key="qwen3.5-27b", | |
| hf="Qwen/Qwen3.5-27B", | |
| revision="fc05daec18b0a78c049392ed2e771dde82bdf654", | |
| n_blocks=64, hidden_size=5120, readout_block=42, | |
| lens=LensFile(repo=LENS_REPO, revision=LENS_REV, | |
| file="qwen3.5-27b/jlens/Salesforce-wikitext/Qwen3.5-27B_jacobian_lens.pt"), | |
| think_end="</think>", aliases=("qwen",), paper_model=True, | |
| )) | |
| # No lens is published for Mistral, so this one was fitted for the paper and mirrored. | |
| _add(ModelSpec( | |
| key="magistral-small-2507", | |
| hf="mistralai/Magistral-Small-2507", | |
| revision="b3583a426ca23186681c4297926dfe9219be6161", | |
| n_blocks=40, hidden_size=5120, readout_block=26, | |
| lens=LensFile(repo="loose-bits/uspace-lenses", | |
| revision="59e7f1c39e6a6379321b55ac395f5d7ed23cebcb", | |
| file="magistral_100fix_jacobian_lens.pt"), | |
| think_end="[/THINK]", think_start="[THINK]", | |
| sampling={"temperature": 0.7, "top_p": 0.95}, | |
| aliases=("magistral",), paper_model=True, | |
| )) | |
| # Not a paper model: small enough to run interactively, and its published lens is fitted on | |
| # more prompts (1000) than any of the three above. | |
| _add(ModelSpec( | |
| key="qwen3.5-4b", | |
| hf="Qwen/Qwen3.5-4B", | |
| revision="851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a", | |
| n_blocks=32, hidden_size=2560, readout_block=20, | |
| lens=LensFile(repo=LENS_REPO, revision=LENS_REV, | |
| file="qwen3.5-4b/jlens/Salesforce-wikitext/Qwen3.5-4B_jacobian_lens_n1000.pt"), | |
| think_end="</think>", aliases=("qwen4b",), | |
| )) | |
| _ALIASES = {a: s.key for s in MODELS.values() for a in s.aliases} | |
| def get(model: str | ModelSpec) -> ModelSpec: | |
| """Look a model up by key or alias.""" | |
| if isinstance(model, ModelSpec): | |
| return model | |
| key = _ALIASES.get(model, model) | |
| if key not in MODELS: | |
| known = ", ".join(sorted(MODELS) + sorted(_ALIASES)) | |
| raise KeyError(f"unknown model {model!r}; known: {known}") | |
| return MODELS[key] | |
| def keys(paper_only: bool = False) -> list[str]: | |
| """Registry keys, optionally only the three models the paper reports.""" | |
| return sorted(k for k, s in MODELS.items() if s.paper_model or not paper_only) | |
| def decode_block(model: str | ModelSpec) -> int: | |
| """The block to read an axis's *identity* at, which is not its readout block.""" | |
| spec = get(model) | |
| return rel_to_block(DECODE_REL, spec.n_blocks) | |