diff --git a/scripts/scratch/calib_pos_C.py b/scripts/scratch/calib_pos_C.py new file mode 100644 index 0000000..6d5cee6 --- /dev/null +++ b/scripts/scratch/calib_pos_C.py @@ -0,0 +1,47 @@ +"""Find the fluent positive-C knee for the word_steering demo prompt. (Claude) +The committed notebook claimed C~0.5 stays fluent and C~1 spams, but at seed=0 the +"describe your week" free-form already degenerates into happy-spam at C=0.5. Sweep +low C to find where the free-form stays coherent, and log the rubric number so we +can pick a C that is BOTH fluent and visibly steered. + + uv run python scripts/scratch/calib_pos_C.py +""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +import config # noqa: E402 +import torch # noqa: E402 +from loguru import logger # noqa: E402 +from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402 + +from jsteer import Jacobian # noqa: E402 +from jsteer.demo import chat_input, rubric_score # 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) +band = jac.steer_band(model) +v = jac.word_vector(model, tok, ["happy", "joy"], layers=band) + +prompt = chat_input(tok, "Describe how your week has been going.") +enc = tok(prompt, return_tensors="pt").to(model.device) +RUBRIC = "On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?" + + +@torch.no_grad() +def gen_tail(C): + torch.manual_seed(0) + with v(model, C=C): + out = model.generate(**enc, max_new_tokens=400, pad_token_id=tok.eos_token_id) + ans, pmass = rubric_score(model, tok, RUBRIC, max_new_tokens=400, seed=0) + txt = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True) + ans_only = txt.split("")[-1].strip().replace("\n", " ")[:220] + return ans_only, ans, pmass + + +for C in (0.0, 0.2, 0.3, 0.4, 0.5): + tail, ans, pmass = gen_tail(C) + logger.info(f"\nC={C:+g} rubric ans={ans:.2f}/9 pmass={pmass:.2f}\n free-form answer: {tail!r}") diff --git a/scripts/scratch/proto_lens_slice.py b/scripts/scratch/proto_lens_slice.py new file mode 100644 index 0000000..16746f4 --- /dev/null +++ b/scripts/scratch/proto_lens_slice.py @@ -0,0 +1,79 @@ +"""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") diff --git a/scripts/scratch/smoke_sweep.py b/scripts/scratch/smoke_sweep.py new file mode 100644 index 0000000..a5f8d70 --- /dev/null +++ b/scripts/scratch/smoke_sweep.py @@ -0,0 +1,35 @@ +"""Smoke: coherence_sweep with sampling on -> ans_std should be >0 (BMA averages over +distinct think traces). (Claude) max_steps=2 keeps it to ~5 C points for speed. + + uv run python scripts/scratch/smoke_sweep.py +""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +import config # 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 jsteer import Jacobian # noqa: E402 +from jsteer.demo import coherence_sweep # 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) +band = jac.steer_band(model) +v = jac.word_vector(model, tok, ["happy", "joy"], layers=band) + +rows = coherence_sweep(model, tok, v, + "On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?", + step=0.1, pmass_floor=0.9, max_steps=2, n_samples=3, max_new_tokens=384) +logger.info("\n" + tabulate(rows, headers="keys", tablefmt="github", floatfmt="+.2f")) +# SHOULD: with do_sample the 3 seeds diverge, so at least one coherent row has ans_std>0. +# If ALL ans_std==0, sampling isn't taking effect (still greedy) -> BMA is a no-op. +nonzero = [r for r in rows if r["ans_std"] > 0] +logger.info(f"\nUAT: rows with ans_std>0 = {len(nonzero)}/{len(rows)} " + f"(SHOULD be >0 -> sampling+BMA active)")