HelioAI's picture
LFM3-1.7B inference modules: experimental release
67ebcef verified
Raw History Blame Contribute Delete
4.25 kB
"""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)