File size: 3,028 Bytes
b9301f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | #!/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()
|