ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | #!/usr/bin/env python3 | |
| """Chat ATOM on the last non-collapsed hard ckpt. No Grok Bot. No extra train. | |
| Preferred checkpoints (first that exists wins): | |
| 1. checkpoints/atom_native_field_next/atom_native.pt @ 3,458,000 logits 0.838 | |
| 2. checkpoints/atom_native_linguistic/ or spans @ 3,425,000 logits 0.866 | |
| 3. checkpoints/atom_native_chat_talk/atom_native.pt | |
| NOT the alpha_local / payload ckpt @ 3,508,000 (logits 0.999 collapse). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| from pathlib import Path | |
| from src.atom_native import AtomNativeModel | |
| CANDIDATES = ( | |
| Path("checkpoints/atom_native_field_next/atom_native.pt"), | |
| Path("checkpoints/atom_native_next_packet/atom_native.pt"), | |
| Path("checkpoints/atom_native_linguistic_spans/atom_native.pt"), | |
| Path("checkpoints/atom_native_linguistic/atom_native.pt"), | |
| Path("checkpoints/atom_native_chat_talk/atom_native.pt"), | |
| Path("checkpoints/atom_native_obligatory_hard/atom_native.pt"), | |
| ) | |
| def resolve_ckpt(explicit: str | None) -> Path: | |
| if explicit: | |
| path = Path(explicit) | |
| if not path.is_file(): | |
| raise SystemExit(f"checkpoint missing: {path}") | |
| return path | |
| for path in CANDIDATES: | |
| if path.is_file() and path.stat().st_size > 0: | |
| return path | |
| raise SystemExit( | |
| "no non-collapsed ckpt found. Pass --checkpoint. " | |
| "Do not use checkpoints/atom_native_alpha_local/" | |
| ) | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="ATOM speech chat (payload-copy OFF)") | |
| parser.add_argument("--checkpoint", default=None) | |
| parser.add_argument("--prompt", default="Bonjour") | |
| parser.add_argument("--max-packets", type=int, default=8) | |
| parser.add_argument("--d-model", type=int, default=64) | |
| parser.add_argument("--n-modes", type=int, default=64) | |
| parser.add_argument("--max-payload-bytes", type=int, default=32) | |
| args = parser.parse_args() | |
| ckpt = resolve_ckpt(args.checkpoint) | |
| print(f"[speech] loading {ckpt}") | |
| model = AtomNativeModel( | |
| d_model=args.d_model, | |
| n_modes=args.n_modes, | |
| n_atoms_max=512, | |
| max_payload_bytes=args.max_payload_bytes, | |
| field_max_rms=3.0, | |
| field_obligatory_hard=True, | |
| enable_merge=True, | |
| ) | |
| extra = model.load(ckpt) or {} | |
| step = extra.get("step") or extra.get("global_step") | |
| print(f"[speech] extra_keys={sorted(extra.keys())[:12]} step={step}") | |
| model.eval() | |
| raw = model.generate_packets( | |
| args.prompt, | |
| max_packets=args.max_packets, | |
| deterministic=True, | |
| prefer_printable=True, | |
| payload_copy=False, | |
| dialogue_wrap=True, | |
| speech_gate=True, | |
| merge_enabled=False, | |
| ) | |
| text = raw.decode("utf-8", errors="replace") | |
| print(f"[speech] prompt={args.prompt!r}") | |
| print(f"[speech] response={text!r}") | |
| if not text.strip(): | |
| print("[speech] empty — speech_ok refused soup. Need the 3.458M ckpt, not 3.508M.") | |
| if __name__ == "__main__": | |
| main() | |