80 lines
3.2 KiB
Python
80 lines
3.2 KiB
Python
"""Companion to probe_discount.py: lens table over the PROMPT positions
|
|
after k=2 prefill settling (arm A adapter) — what the workspace holds
|
|
about the question before any pause/answer compute begins."""
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from carry_common import _pos_ids, prompt_prefill
|
|
from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from jlens.core import JLens, ResidualCapture, load_model # noqa: E402
|
|
|
|
OUT = Path(os.environ.get("LOOP_OUT",
|
|
Path(__file__).resolve().parent.parent / "results-loop"))
|
|
Q = "A shirt costs $80 and is on sale for 15% off. How much does it cost?"
|
|
K = 2
|
|
TAG = "discount"
|
|
if "--idx" in sys.argv:
|
|
import json
|
|
_i = int(sys.argv[sys.argv.index("--idx") + 1])
|
|
_star = {it["idx"]: it for it in json.load(open(OUT / "star_data.json"))
|
|
if it["split"] == "test"}
|
|
Q = _star[_i]["question"]
|
|
TAG = f"gsm{_i}"
|
|
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
looper = BandLooper(model)
|
|
band = list(range(looper.l0, looper.l1 + 1))
|
|
adapter = MergeAdapter(d=model.config.get_text_config().hidden_size).cuda()
|
|
adapter.load_state_dict(torch.load(OUT / "adapter_carrycot_e400.pt",
|
|
map_location="cuda"))
|
|
jbar = torch.load(Path(__file__).resolve().parent.parent / "results/jbar.pt",
|
|
map_location="cuda")["Jbar"]
|
|
jl = JLens(model, tok, jbar)
|
|
|
|
prompt = chat_prompt(tok, Q, DIRECT_SUFFIX)
|
|
ids = tok(prompt, add_special_tokens=False,
|
|
return_tensors="pt")["input_ids"].cuda()
|
|
mask = torch.ones_like(ids)
|
|
|
|
with torch.no_grad():
|
|
calls, _ = looper.capture(ids, mask, logits_to_keep=1,
|
|
position_ids=_pos_ids(mask))
|
|
e = looper._hin[looper.l0]
|
|
pmask = torch.ones_like(ids, dtype=torch.bool)
|
|
S, X = prompt_prefill(looper, adapter, e, calls, pmask, K)
|
|
with ResidualCapture(model, layers=band) as rc:
|
|
looper.band(X, calls)
|
|
acts = {l: rc.acts[l][0] for l in band}
|
|
# k=0 reference: plain forward, no adapter anywhere
|
|
with ResidualCapture(model, layers=band) as rc0:
|
|
model(input_ids=ids, attention_mask=mask, use_cache=False)
|
|
acts0 = {l: rc0.acts[l][0] for l in band}
|
|
|
|
table = {}
|
|
for l in band:
|
|
idx, p = jl.read(acts[l], l, topk=5)
|
|
table[l] = (idx.cpu(), p.cpu())
|
|
idx0, p0 = jl.read(acts0[30], 30, topk=5)
|
|
toks = [repr(tok.decode([t])) for t in ids[0].tolist()]
|
|
torch.save({"question": Q, "k": K, "poslab": toks, "band": band,
|
|
"table": table, "base30": (idx0.cpu(), p0.cpu())},
|
|
OUT / f"probe_{TAG}_prompt.pt")
|
|
show = [14, 18, 22, 26, 30]
|
|
print(f"=== prompt board after k={K} settle (arm A) || k=0 base ===")
|
|
for i, lab in enumerate(toks):
|
|
row, row0 = [], []
|
|
for l in show:
|
|
idx, _ = jl.read(acts[l][i][None], l, topk=3)
|
|
row.append("/".join(repr(tok.decode([t]))[1:-1]
|
|
for t in idx[0].tolist()))
|
|
idx0, _ = jl.read(acts0[30][i][None], 30, topk=3)
|
|
base30 = "/".join(repr(tok.decode([t]))[1:-1] for t in idx0[0].tolist())
|
|
print(f"{i:3d} {lab:14s} " + " ".join(
|
|
f"L{l}:{r}" for l, r in zip(show, row)) + f" || base-L30:{base30}",
|
|
flush=True)
|