ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 3)
969c687 verified | """(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" | |
| 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 | |
| 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() | |