"""(2) Train/generation consistency of ATOM @4996k. A. teacher-forced path (Atomizer.encode of whole text, forward_packet per byte == eval_windows/transition_loss) vs the generate_packets path with the continuation FORCED (same calls as src/atom_native.py:2258-2307, i.e. prompt encode + packet_from_payload(gold byte)). Compare raw logits per step + packet features. B. same with packet_from_payload boundary patched to the train-time label (diagnostic monkeypatch only). C. chat samples (exact probe settings) original vs patched, doubled-punctuation rate. """ import json, sys, re from pathlib import Path import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) from common import CKPT, PROMPTS, load_model from src.io.atomizer import Atomizer from src.speech_lock import format_dialogue_prompt OUT = Path(__file__).resolve().parent torch.set_num_threads(1) CASES = [("Bonjour, comment ça va ?", "Ça va très bien, et toi ?\n"), ("Qui es-tu ?", "Je suis Atom, un petit modèle qui apprend à parler.\n"), ("Il était une fois", "Il était une fois un petit robot qui apprenait des mots.\n"), ("Tu habites où ?", "J'habite à Strasbourg. Et toi ?\n")] ORIG_INFER = Atomizer._infer_commit_boundary.__func__ def train_like_boundary(cls, payload): # what encode() emits for a 1-byte packet in byte-tick return "newline" if payload[-1] in (10, 13) else "max_span" @torch.no_grad() def tf_logits(m, text): m.eval(); m.reset_state(reset_atomizer=True) pk = m.atomizer.encode(text, reset=True) return [m.forward_packet(p)["surface"]["byte_logits"][0].clone() for p in pk[:-1]], pk @torch.no_grad() def gen_forced_logits(m, prompt, gold: bytes): m.eval(); m.surface.set_obligatory_hard(True); m.surface.payload_enabled = False m.reset_state(reset_atomizer=True) wrapped = format_dialogue_prompt(prompt) pp = m.atomizer.encode(wrapped, reset=True) logits, pkts = [], list(pp) for p in pp[:-1]: logits.append(m.forward_packet(p)["surface"]["byte_logits"][0].clone()) cur = pp[-1] for b in gold: logits.append(m.forward_packet(cur)["surface"]["byte_logits"][0].clone()) cur = m.atomizer.packet_from_payload(bytes([b])); pkts.append(cur) return logits, wrapped, pkts def compare(m, tag): rows = [] for prompt, gold in CASES: g = gold.encode() lg_g, wrapped, pk_g = gen_forced_logits(m, prompt, g) full = wrapped + gold lg_t, pk_t = tf_logits(m, full) fb = full.encode(); n0 = len(wrapped.encode()) n = min(len(lg_t), len(lg_g)) diffs = [float((lg_t[i] - lg_g[i]).abs().max()) for i in range(n)] ce_t = [float(-torch.log_softmax(lg_t[i], -1)[fb[i + 1]]) for i in range(n)] ce_g = [float(-torch.log_softmax(lg_g[i], -1)[fb[i + 1]]) for i in range(n)] feat = [float((pk_t[i].features - pk_g[i].features).abs().max()) for i in range(min(len(pk_t), len(pk_g)))] bad = [i for i in range(n) if diffs[i] > 1e-4] first = bad[0] if bad else None rows.append({ "prompt": prompt, "gold": gold, "n_steps": n, "n_prompt_bytes": n0, "max_logit_diff_prompt_part": max(diffs[:n0 - 1]) if n0 > 1 else 0.0, "max_logit_diff_reply_part": max(diffs[n0 - 1:]), "n_mismatch_steps(>1e-4)": len(bad), "first_mismatch_step": first, "first_mismatch_input_byte": (chr(fb[first]) if first is not None else None), "packet_feature_mismatch_bytes": [chr(fb[i]) for i, d in enumerate(feat) if d > 1e-6], "boundaries_tf_vs_gen_at_mismatch": [(chr(fb[i]), pk_t[i].boundary, pk_g[i].boundary) for i, d in enumerate(feat) if d > 1e-6][:12], "reply_ce_teacher_forced": sum(ce_t[n0 - 1:]) / len(ce_t[n0 - 1:]), "reply_ce_gen_path_forced": sum(ce_g[n0 - 1:]) / len(ce_g[n0 - 1:]), "per_step": [(chr(fb[i]) if fb[i] < 128 else f"0x{fb[i]:02x}", round(diffs[i], 4), round(ce_t[i], 3), round(ce_g[i], 3), # top-3 next byte under gen path [chr(int(j)) if int(j) < 128 else f"0x{int(j):02x}" for j in torch.topk(lg_g[i], 3).indices]) for i in range(n0 - 1, n)], }) print(tag, json.dumps([{k: v for k, v in r.items() if k != "per_step"} for r in rows], ensure_ascii=False, indent=1), flush=True) return rows def chat_samples(m, seeds=range(8)): out = {} for p in PROMPTS: out[p] = [] for s in seeds: torch.manual_seed(s) raw = m.generate_packets(p, max_packets=96, temperature=0.8, top_k=12, deterministic=False, max_length=100, prefer_printable=True, reset=True, merge_enabled=False, speech_gate=True) out[p].append(raw.decode("utf-8", "replace")) return out def chat_samples_nogate(m, seeds=range(8)): out = {} for p in PROMPTS: out[p] = [] for s in seeds: torch.manual_seed(s) raw = m.generate_packets(p, max_packets=96, temperature=0.8, top_k=12, deterministic=False, max_length=100, prefer_printable=True, reset=True, merge_enabled=False, speech_gate=False, anti_repeat=False) out[p].append(raw.decode("utf-8", "replace")) return out def doubled_rate(samples): txt = "".join(s for v in samples.values() for s in v) punct = sum(1 for c in txt if c in ",.-'?!:;") dbl = len(re.findall(r"([,.\-'?!:;])\1", txt)) return {"chars": len(txt), "punct": punct, "doubled_punct_pairs": dbl} def main(): m = load_model(CKPT) res = {"ckpt": CKPT} res["A_original"] = compare(m, "A original") res["C_chat_original"] = chat_samples(m); res["C_rate_original"] = doubled_rate(res["C_chat_original"]) res["C_chat_original_nogate"] = chat_samples_nogate(m); res["C_rate_original_nogate"] = doubled_rate(res["C_chat_original_nogate"]) print(json.dumps(res["C_chat_original"], ensure_ascii=False, indent=1), res["C_rate_original"], flush=True) print(json.dumps(res["C_chat_original_nogate"], ensure_ascii=False, indent=1), res["C_rate_original_nogate"], flush=True) Atomizer._infer_commit_boundary = classmethod(train_like_boundary) try: res["B_patched"] = compare(m, "B patched") res["C_chat_patched"] = chat_samples(m); res["C_rate_patched"] = doubled_rate(res["C_chat_patched"]) res["C_chat_patched_nogate"] = chat_samples_nogate(m); res["C_rate_patched_nogate"] = doubled_rate(res["C_chat_patched_nogate"]) print(json.dumps(res["C_chat_patched"], ensure_ascii=False, indent=1), res["C_rate_patched"], flush=True) print(json.dumps(res["C_chat_patched_nogate"], ensure_ascii=False, indent=1), res["C_rate_patched_nogate"], flush=True) finally: Atomizer._infer_commit_boundary = classmethod(ORIG_INFER) (OUT / "decode_consistency.json").write_text(json.dumps(res, ensure_ascii=False, indent=1)) if __name__ == "__main__": main()