| """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 | |
| 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() | |