mirror of
https://github.com/wassname/jsteer.git
synced 2026-09-09 11:25:03 +08:00
feat: load authors' pre-fitted n1000 Hub lenses; steer_band + masked lens_topk
config.HUB_LENS_FILE maps HF model -> the authors' pre-fitted Jacobian lens on the Hub (neuronpedia/jacobian-lens, raw Salesforce-wikitext, n=1000). Loading one beats fitting locally: same estimator, 1000 prompts, zero compute. Our Jacobian already wraps jlens.JacobianLens, so their .pt loads through Jacobian.from_pretrained with no format change (verified: n1000 4B loads, d_model=2560, layers [0..30]). jacobian.py: - steer_band(model, lo=0.3, hi=0.9): pre-fitted lenses span every layer; steering all of them over-drives the residual, so restrict to the mid-depth band run-524 used. - lens_topk reuses jlens.vis._meaningful_token_mask so j-space readouts hide punctuation/single-char/special tokens (per the walkthrough these trail the interesting word tokens on Qwen). Verified: Eiffel Tower resolves city->Paris clean. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
@@ -26,6 +26,27 @@ ART = ROOT / "artifacts"
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
|
||||
# The authors publish pre-fitted Jacobian lenses on the Hub (raw Salesforce-
|
||||
# wikitext, n=1000 where the _n1000 suffix is present). Loading one beats fitting
|
||||
# locally: identical estimator, 1000 prompts, zero compute. Keyed by HF model id;
|
||||
# see github.com/anthropics/jacobian-lens walkthrough.ipynb.
|
||||
LENS_REPO = "neuronpedia/jacobian-lens"
|
||||
LENS_REVISION = "qwen-n1000"
|
||||
HUB_LENS_FILE = {
|
||||
"Qwen/Qwen3.5-4B": "qwen3.5-4b/jlens/Salesforce-wikitext/Qwen3.5-4B_jacobian_lens_n1000.pt",
|
||||
"Qwen/Qwen3.6-27B": "qwen3.6-27b/jlens/Salesforce-wikitext/Qwen3.6-27B_jacobian_lens_n1000.pt",
|
||||
"Qwen/Qwen3-4B": "qwen3-4b/jlens/Salesforce-wikitext/Qwen3-4B_jacobian_lens.pt",
|
||||
"Qwen/Qwen3-8B": "qwen3-8b/jlens/Salesforce-wikitext/Qwen3-8B_jacobian_lens.pt",
|
||||
"Qwen/Qwen3-14B": "qwen3-14b/jlens/Salesforce-wikitext/Qwen3-14B_jacobian_lens.pt",
|
||||
"Qwen/Qwen3-32B": "qwen3-32b/jlens/Salesforce-wikitext/Qwen3-32B_jacobian_lens.pt",
|
||||
}
|
||||
|
||||
|
||||
def hub_lens_file(model_name: str) -> str:
|
||||
"""Filename of the authors' pre-fitted lens for `model_name` inside LENS_REPO.
|
||||
KeyError (fail fast) if they don't publish one -- then fit locally via fit.py."""
|
||||
return HUB_LENS_FILE[model_name]
|
||||
|
||||
|
||||
def chat_corpus(tok, n_prompts: int) -> list[str]:
|
||||
"""jlens's WikiText prompts wrapped in the chat template. Fitting on chat-
|
||||
|
||||
+19
-3
@@ -226,6 +226,13 @@ class Jacobian:
|
||||
" ".join(f"{l}:{per_layer[l].norm():.3g}" for l in cfg.layers))
|
||||
return _to_vector(cfg, per_layer)
|
||||
|
||||
def steer_band(self, model, *, lo: float = 0.3, hi: float = 0.9) -> tuple[int, ...]:
|
||||
"""Fitted layers within the [lo, hi] fraction of model depth. The authors'
|
||||
pre-fitted lenses span EVERY layer; steering all of them at once over-drives
|
||||
the residual, so restrict to the mid-depth band run-524 used."""
|
||||
n = model.config.num_hidden_layers
|
||||
return tuple(l for l in self.lens.source_layers if lo <= l / n <= hi)
|
||||
|
||||
def _steer_layers(self, layers) -> tuple[int, ...]:
|
||||
"""None -> all fitted layers; else explicit int indices (float bands are
|
||||
a fit-time concept -- the lens doesn't know n_layers to resolve them)."""
|
||||
@@ -297,13 +304,22 @@ class Jacobian:
|
||||
# -- bonus: the lens's native forward readout --------------------------------
|
||||
|
||||
def lens_topk(self, model, tok, prompt: str, layer: int, *, k: int = 10,
|
||||
position: int = -1) -> list[tuple[str, float]]:
|
||||
position: int = -1, mask_wordlike: bool = True) -> list[tuple[str, float]]:
|
||||
"""Lens readout at `layer`: transport the residual to the final basis
|
||||
with J_l and decode to tokens (a linear approximation, not the literal
|
||||
computation). jlens's native use, handy in demos."""
|
||||
computation). jlens's native use, handy in demos.
|
||||
|
||||
`mask_wordlike` reuses jlens's own word-like vocab mask so the readout
|
||||
hides punctuation/single-char/special tokens (which, per the walkthrough,
|
||||
trail the interesting word tokens on Qwen); ranks are unaffected."""
|
||||
from jlens.vis import _meaningful_token_mask
|
||||
lm = from_hf(model, tok)
|
||||
lens_logits, _, _ = self.lens.apply(lm, prompt, layers=[layer],
|
||||
positions=[position])
|
||||
top = lens_logits[layer][0].topk(k)
|
||||
logits = lens_logits[layer][0]
|
||||
if mask_wordlike:
|
||||
wl = _meaningful_token_mask(tok, logits.shape[-1], logits.device)
|
||||
logits = logits.masked_fill(~wl, float("-inf"))
|
||||
top = logits.topk(k)
|
||||
return [(tok.decode([i]), float(v)) for i, v in
|
||||
zip(top.indices.tolist(), top.values.tolist())]
|
||||
|
||||
Reference in New Issue
Block a user