File size: 4,253 Bytes
67ebcef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """Steering direction 'hesitation' from the model's own reasoning traces (no labels, no external model).
At every paragraph break ("\n\n") inside <think>, the residual stream of each layer is recorded at the break token,
i.e. at the moment the model decides how to continue. Paragraphs that open with a hesitation/re-check marker
("Wait", "Alternatively", "Hmm", "Let me double-check", ...) form class R, all others class E.
direction[layer] = mean_R - mean_E. Saved as .pt (all layers) for building llama.cpp control vectors.
usage: extract_cvec.py --gen results/dev_orig/dev.jsonl --items data/train/dev.jsonl --out results/cvec/hes.pt
"""
import argparse, json, re, torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import transformers.integrations.sdpa_attention as _sdpa
# Windows torch has no flash kernel: masked/GQA SDPA would fall back to the O(L^2)-memory math kernel.
# Expanding K/V instead keeps it on the memory-efficient kernel.
_sdpa.use_gqa_in_sdpa = lambda *a, **k: False
ap = argparse.ArgumentParser()
ap.add_argument("--gen", required=True)
ap.add_argument("--items", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--model", default="models/LFM3-1.7B")
ap.add_argument("--max-resp", type=int, default=200)
ap.add_argument("--max-len", type=int, default=8192)
args = ap.parse_args()
HES = re.compile(r"^(\*\*)?(wait|but wait|hmm|alternatively|actually|hold on|oh wait|however, wait|let me (double[- ]check|"
r"check (again|this|that)|verify|re-?examine|reconsider|think again|re-?check|re-?read|make sure)|"
r"but let me|let's (double[- ]check|verify|check again))", re.I)
tok = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16, attn_implementation="sdpa", trust_remote_code=True).cuda().eval()
L = model.config.num_hidden_layers
items = {json.loads(l)["id"]: json.loads(l) for l in open(args.items, encoding="utf-8")}
recs = [json.loads(l) for l in open(args.gen, encoding="utf-8")]
recs = [r for r in recs if items[r["id"]]["task"] == "math" and "tokens" in r][: args.max_resp]
cap = {}
hooks = [layer.register_forward_hook(lambda m, i, o, k=k: cap.__setitem__(k, (o[0] if isinstance(o, tuple) else o)[0]))
for k, layer in enumerate(model.model.layers)]
sums = {c: torch.zeros(L, model.config.hidden_size, dtype=torch.float64) for c in "RE"}
cnt = {"R": 0, "E": 0}
norm = torch.zeros(L, dtype=torch.float64)
for n_r, r in enumerate(recs):
p = tok.apply_chat_template([{"role": "user", "content": items[r["id"]]["prompt"]}], tokenize=True,
add_generation_prompt=True, return_dict=False)
resp = r["tokens"]
pieces = tok.convert_ids_to_tokens(resp)
pieces = [tok.convert_tokens_to_string([x]) for x in pieces]
text, pos, labels = "", [], []
for i, s in enumerate(pieces):
text += s
if "</think>" in text:
break
if s.endswith("\n\n") and i + 1 < len(resp):
nxt = tok.decode(resp[i + 1: i + 12]).lstrip()
pos.append(len(p) + i)
labels.append("R" if HES.match(nxt) else "E")
ids = (p + resp)[: args.max_len]
keep = [(q, c) for q, c in zip(pos, labels) if q < len(ids)]
if not keep:
continue
with torch.no_grad():
model(input_ids=torch.tensor([ids], device="cuda"), use_cache=False)
qi = torch.tensor([q for q, _ in keep], device="cuda")
isr = torch.tensor([c == "R" for _, c in keep])
for k in range(L):
h = cap[k][qi].double().cpu()
sums["R"][k] += h[isr].sum(0)
sums["E"][k] += h[~isr].sum(0)
norm[k] += h.norm(dim=-1).sum()
for _, c in keep:
cnt[c] += 1
if n_r % 20 == 0:
print(n_r, cnt, flush=True)
mR, mE = sums["R"] / cnt["R"], sums["E"] / cnt["E"]
d = (mR - mE).float()
avg_norm = (norm / (cnt["R"] + cnt["E"])).float()
for k in range(L):
print(f"layer {k:2d} |dir| {d[k].norm():8.2f} |h| {avg_norm[k]:8.2f} ratio {d[k].norm() / avg_norm[k]:.3f} "
f"cos(mR,mE) {torch.nn.functional.cosine_similarity(mR[k], mE[k], 0):.3f}")
torch.save({"dir": d, "avg_norm": avg_norm, "count": cnt}, args.out)
print("saved", args.out, cnt)
|