"""Does the model follow a subject change, or keep answering the old one? Score is keyword hits on the new topic over hits on both: adherence = hits(B) / (hits(A) + hits(B)) 1.0 switched, 0.0 stuck. Depth padding repeats the last A turn so a longer history can pull harder. Training topics live in anti_echo.py, not here. """ import argparse import json import os import re import sys os.environ.setdefault("TRANSFORMERS_NO_TF", "1") os.environ.setdefault("USE_TF", "0") import torch from transformers import AutoModelForCausalLM sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from tokenizer_shim import RawTok CHATML_U = "<|im_start|>user\n{}<|im_end|>\n" CHATML_A = "<|im_start|>assistant\n{}<|im_end|>\n" OPEN_A = "<|im_start|>assistant\n" STOPS = ["<|im_end|>", "<|im_start|>"] CASES = [ ([("Tell me about the steam engine.", "The steam engine converts heat into mechanical work using pressurised steam."), ("Who improved its efficiency?", "James Watt, with the separate condenser.")], "What is a good recipe for bread?", ["steam", "engine", "watt", "condenser", "heat", "piston", "boiler", "coal"], ["bread", "flour", "dough", "yeast", "bake", "oven", "loaf", "knead", "recipe"]), ([("How do consensus protocols work?", "They let machines agree on a value despite failures and delays."), ("Which ones are common?", "Paxos and Raft are the best known.")], "What should I feed a young rabbit?", ["consensus", "protocol", "paxos", "raft", "machine", "node", "leader", "log"], ["rabbit", "hay", "pellet", "feed", "vegetable", "water", "diet", "carrot"]), ([("Explain photosynthesis.", "Plants convert light energy into chemical energy stored as sugars."), ("What role does chlorophyll play?", "It absorbs light, mostly in the red and blue parts of the spectrum.")], "How do I change a car tyre?", ["photosynthesis", "chlorophyll", "light", "plant", "sugar", "leaf", "carbon"], ["tyre", "tire", "car", "jack", "wheel", "bolt", "spare", "wrench", "nut"]), ([("What caused the Great Depression?", "A mix of a stock market crash, bank failures and falling demand."), ("How long did it last?", "Roughly through the 1930s in most countries.")], "What are the rules of chess?", ["depression", "economy", "bank", "market", "crash", "unemployment", "1930"], ["chess", "piece", "pawn", "king", "queen", "board", "check", "move", "rook"]), ([("Describe how rain forms.", "Water vapour condenses around particles until droplets grow heavy enough to fall."), ("What is a cold front?", "A boundary where colder air pushes under warmer air.")], "How do I learn to play guitar?", ["rain", "cloud", "vapour", "vapor", "front", "air", "weather", "droplet", "storm"], ["guitar", "chord", "string", "fret", "play", "practice", "strum", "music"]), ([("What is a black hole?", "A region where gravity is so strong that nothing escapes, not even light."), ("What is the event horizon?", "The boundary beyond which escape is impossible.")], "What should I plant in a vegetable garden?", ["black", "hole", "gravity", "horizon", "star", "light", "mass", "space"], ["plant", "garden", "vegetable", "seed", "soil", "grow", "tomato", "harvest"]), ] def build(turns, question, depth_pad=0): parts = [] for u, a in turns: parts.append(CHATML_U.format(u)) parts.append(CHATML_A.format(a)) for _ in range(depth_pad): parts.append(CHATML_U.format(turns[-1][0])) parts.append(CHATML_A.format(turns[-1][1])) parts.append(CHATML_U.format(question)) parts.append(OPEN_A) return "".join(parts) def _cut_at_stop(txt): cuts = [txt.find(s) for s in STOPS if txt.find(s) >= 0] return txt[:min(cuts)] if cuts else txt @torch.no_grad() def gen(model, tok, dev, ctx, max_new=110): ids = tok(ctx).input_ids.to(dev) past = 0 cur = ids cache = None out = [] for _ in range(max_new): pos = torch.arange(past, past + cur.shape[1], device=dev).unsqueeze(0) r = model(cur, position_ids=pos, past_key_values=cache, use_cache=True) nxt = int(r.logits[0, -1].argmax()) out.append(nxt) past += cur.shape[1] cur = torch.tensor([[nxt]], device=dev) cache = r.past_key_values return _cut_at_stop(tok.decode(out)) def hits(text, kws): bag = set(re.findall(r"[a-z]+", text.lower())) return sum(1 for k in kws if k in bag) def main(): ap = argparse.ArgumentParser() ap.add_argument("--models", required=True, help="comma list of name=path") ap.add_argument("--device", default="cuda") ap.add_argument("--depths", default="0,2") ap.add_argument("--out", default=None) a = ap.parse_args() depths = [int(d) for d in a.depths.split(",")] rows = {} dump = [] for spec in a.models.split(","): name, path = spec.split("=", 1) tok = RawTok(path) m = AutoModelForCausalLM.from_pretrained( path, trust_remote_code=True, dtype=torch.float32 ).to(a.device).eval() rows[name] = {} for d in depths: adh = [] stuck = 0 for turns, q, kw_a, kw_b in CASES: txt = gen(m, tok, a.device, build(turns, q, d)) ha, hb = hits(txt, kw_a), hits(txt, kw_b) score = hb / (ha + hb) if (ha + hb) else 0.5 adh.append(score) stuck += (ha > 0 and hb == 0) if d == depths[0]: dump.append({ "model": name, "q": q, "old_hits": ha, "new_hits": hb, "text": txt[:260], }) rows[name][d] = (sum(adh) / len(adh), stuck, len(CASES)) del m if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"{'model':<22}" + "".join(f"{'depth+' + str(d):>16}" for d in depths)) print(f"{'':<22}" + "".join(f"{'adher stuck':>16}" for _ in depths)) for name, per in rows.items(): line = f"{name:<22}" for d in depths: s, st, n = per[d] line += f"{s:>9.2f} {st}/{n:>4}" print(line) print("\nadherence 1.0 = fully switched to new topic, 0.0 = stuck on old") print("stuck = replies using ONLY old-topic vocabulary") if a.out: payload = { "rows": {k: {str(d): v for d, v in p.items()} for k, p in rows.items()}, "samples": dump, } json.dump(payload, open(a.out, "w"), indent=2) txt_path = a.out.replace(".json", ".txt") with open(txt_path, "w", encoding="utf-8") as fh: for r in dump: fh.write( f"=== {r['model']} old={r['old_hits']} new={r['new_hits']}\n" f"Q: {r['q']}\n{r['text']!r}\n\n" ) print(f"wrote {a.out} and .txt") if __name__ == "__main__": main()