mirror of
https://github.com/wassname/jsteer.git
synced 2026-09-09 11:25:03 +08:00
wassnameandClaudypoo
339c67fefe
calib_pos_C (fluent +C knee), smoke_sweep (coherence_sweep sampling gives ans_std>0), proto_lens_slice (reference compute_slice readout prototype: auto-tracked tokens incl 巴黎, Paris rank->0 at the model row). Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
80 lines
3.7 KiB
Python
80 lines
3.7 KiB
Python
"""Prototype: use the REFERENCE jlens.vis machinery (compute_slice) for the
|
|
j-lens readout instead of our hand-picked lens_rank. (Claude)
|
|
|
|
compute_slice does the "complex logic over layers and tokens" wassname remembered:
|
|
auto-selects tracked tokens by a frequency-weighted 1/(rank+1) score over the whole
|
|
top-N grid, sweeps every fitted layer, appends the final layer as the J=I model row,
|
|
and returns full-vocab rank tensors. We pin the answer token, window to the last
|
|
position, and pull rank_tensor into a table + a rank-vs-layer plot -- reference logic,
|
|
our presentation, no fork.
|
|
|
|
uv run python scripts/scratch/proto_lens_slice.py
|
|
|
|
SHOULD: ' Paris' rank collapses toward 0 (top) at the deepest layers and the final
|
|
(J=I / model) row; generic ' city' is low-rank mid-depth then climbs; the auto-tracked
|
|
set contains Paris/city without us hand-listing every token. ELSE compute_slice wiring
|
|
or token pinning is wrong.
|
|
"""
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
import matplotlib # noqa: E402
|
|
matplotlib.use("Agg")
|
|
import config # noqa: E402
|
|
import matplotlib.pyplot as plt # noqa: E402
|
|
import torch # noqa: E402
|
|
from loguru import logger # noqa: E402
|
|
from tabulate import tabulate # noqa: E402
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402
|
|
|
|
from jlens.hf import from_hf # noqa: E402
|
|
from jlens.vis import build_page, compute_slice, notebook_iframe # noqa: E402
|
|
from jsteer import Jacobian # noqa: E402
|
|
|
|
MODEL = "Qwen/Qwen3.5-4B"
|
|
tok = AutoTokenizer.from_pretrained(MODEL)
|
|
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("cuda").eval()
|
|
jac = Jacobian.from_pretrained(config.LENS_REPO, filename=config.hub_lens_file(MODEL),
|
|
revision=config.LENS_REVISION)
|
|
|
|
PROMPT = "The Eiffel Tower is located in the city of"
|
|
lm = from_hf(model, tok)
|
|
pin = {tok(" Paris", add_special_tokens=False).input_ids[0],
|
|
tok(" city", add_special_tokens=False).input_ids[0]}
|
|
sd = compute_slice(lm, jac.lens, PROMPT, top_n=10, max_tracked=6,
|
|
pinned_token_ids=pin, last_n_tokens=1, mask_display=True)
|
|
|
|
logger.info(f"layers={sd.layers}")
|
|
logger.info(f"tracked={[sd.vocab_fragment[t] for t in sd.tracked_token_ids]}")
|
|
|
|
# rank of each tracked token across layers at the (single) last position
|
|
ranks = sd.rank_tensor[-1] # [n_layers, n_tracked]
|
|
labels = [sd.vocab_fragment[t] for t in sd.tracked_token_ids]
|
|
# table: show a few layers to stay narrow (first, mid, last-fitted, model row = last)
|
|
show = [0, len(sd.layers) // 2, len(sd.layers) - 1]
|
|
rows = [[labels[j]] + [int(ranks[i, j]) for i in show] for j in range(len(labels))]
|
|
hdr = ["token"] + [f"L{sd.layers[i]}" + ("(model)" if i == len(sd.layers) - 1 else "")
|
|
for i in show]
|
|
logger.info("\n" + tabulate(rows, headers=hdr, tablefmt="github"))
|
|
|
|
# plot: rank+1 (log, inverted so top-of-plot = rank 0 = the model's next token)
|
|
fig, ax = plt.subplots(figsize=(6, 3.2))
|
|
for j, lab in enumerate(labels):
|
|
ax.plot(sd.layers, ranks[:, j] + 1, marker="o", ms=3, label=repr(lab))
|
|
ax.set_yscale("log")
|
|
ax.invert_yaxis()
|
|
ax.axvline(sd.layers[-1], color="0.85", lw=0.8, zorder=0) # J=I model row
|
|
ax.set_xlabel("layer (rightmost = final, J=I = model)")
|
|
ax.set_ylabel("rank+1 (log; 1 = top token)")
|
|
ax.set_title("lens rank of tracked tokens vs depth")
|
|
ax.legend(fontsize=7, ncol=2)
|
|
fig.tight_layout()
|
|
fig.savefig("/tmp/claude-1000/proto_lens_slice.png", dpi=110)
|
|
logger.info("wrote /tmp/claude-1000/proto_lens_slice.png")
|
|
|
|
# and confirm the reference's own HTML page builds without error (heatmap view)
|
|
page, w, h = build_page(sd, PROMPT, title="lens slice", description="proto")
|
|
_ = notebook_iframe(page)
|
|
logger.info(f"build_page OK: {w}x{h} grid, page {len(page)} chars; notebook_iframe OK")
|