"""Chat with an atom-native checkpoint (Atomizer encode → generate → UTF-8). Usage: PYTHONPATH=. python tools/chat_atom_native.py \ --checkpoint checkpoints/atom_native_chat/atom_native.pt \ --prompt "Bonjour" PYTHONPATH=. python tools/chat_atom_native.py \ --checkpoint checkpoints/atom_native_chat/atom_native.pt \ --interactive """ from __future__ import annotations import argparse import sys from pathlib import Path import torch from src.atom_native import DEFAULT_ENERGY_DECAY_BOUNDS, AtomNativeModel def load_model(checkpoint: str | Path, device: str = "cpu") -> AtomNativeModel: path = Path(checkpoint) if not path.exists(): raise FileNotFoundError(f"checkpoint not found: {path}") blob = torch.load(path, map_location="cpu", weights_only=False) config = blob.get("config") or {} bounds = config.get("energy_decay_bounds") energy_decay_bounds = tuple(bounds) if bounds is not None else DEFAULT_ENERGY_DECAY_BOUNDS # Prefer checkpoint field_max_rms; fall back to current chat default (3.0). field_max_rms = config.get("field_max_rms") if field_max_rms is None: field_max_rms = 3.0 hard = bool(config.get("field_obligatory_hard", False)) model = AtomNativeModel( d_model=int(config.get("d_model", 64)), n_modes=int(config.get("n_modes", 64)), n_atoms_max=int(config.get("n_atoms_max", 512)), max_payload_bytes=int(config.get("max_payload_bytes", 16)), field_max_rms=field_max_rms, energy_decay_bounds=energy_decay_bounds, field_obligatory_hard=hard, field_obligatory_readout=bool(config.get("field_obligatory_readout", False)) or hard, last_atom_readout=bool(config.get("last_atom_readout", False)), ) training = model.load(path) if "last_atom_scale" in config: model.surface.last_atom_scale = float(config["last_atom_scale"]) if "last_atom_scale_adaptive" in config: model.surface.last_atom_scale_adaptive = bool(config["last_atom_scale_adaptive"]) if "last_atom_scale_min" in config: model.surface.last_atom_scale_min = float(config["last_atom_scale_min"]) if "last_atom_scale_max" in config: model.surface.last_atom_scale_max = float(config["last_atom_scale_max"]) if hard or bool(config.get("field_obligatory_hard", False)): model.field_obligatory_hard = True model.field_obligatory_readout = True model.surface.set_obligatory_hard(True) # Stay on hard α-MLP; prefer_printable is forced in generate_packets. assert model.surface.field_obligatory_hard, "chat must stay on hard α-MLP path" dyn = (training or {}).get("dynamics_on_load") or {} if dyn.get("energy_decay_repaired"): print( f"[load] repaired energy_decay " f"{dyn.get('energy_decay_before')} -> {dyn.get('energy_decay')} " f"bounds={dyn.get('energy_decay_bounds')}" ) model.to(device) model.eval() return model def generate_reply( model: AtomNativeModel, prompt: str, max_packets: int = 24, temperature: float = 0.7, top_k: int = 8, max_length: int | None = 256, deterministic: bool = False, role_prime: bool = True, merge_enabled: bool = False, ) -> str: # Dialogue-style priming: if the user did not already include a role tag, # wrap as Utilisateur/Assistant so the field sees the training pattern. # MERGE stays off on chat ingest (train still uses enable_merge thr=0.45). primed = prompt if role_prime and "Assistant:" not in prompt and "Utilisateur:" not in prompt: primed = f"Utilisateur: {prompt}\nAssistant:" raw = model.generate_packets( primed, max_packets=max_packets, temperature=temperature, top_k=top_k, deterministic=deterministic, max_length=max_length, prefer_printable=True, reset=True, merge_enabled=merge_enabled, speech_gate=True, ) text = raw.decode("utf-8", errors="replace") # Drop leading whitespace-only noise common before content stabilizes. return text.lstrip("\n\r ") def interactive_loop(model: AtomNativeModel, args: argparse.Namespace) -> None: print("ATOM-native chat (Atomizer + toroidal field). Commands: quit | status | save") while True: try: user = input("\nVous: ").strip() except (EOFError, KeyboardInterrupt): print() break if not user: continue lower = user.lower() if lower in {"quit", "exit", "q"}: break if lower == "status": print( f"atoms={len(model.core.atoms)} " f"d_model={model.core.encoder.d_model} " f"n_modes={model.core.state.n_modes} " f"max_payload={model.max_payload_bytes}" ) continue if lower == "save": out = Path(args.checkpoint).with_name("interactive_atom_native.pt") model.save(out) print(f"saved {out}") continue reply = generate_reply( model, user, max_packets=args.max_packets, temperature=args.temperature, top_k=args.top_k, max_length=args.max_length, deterministic=args.deterministic, role_prime=not args.no_role_prime, merge_enabled=bool(args.merge_ingest), ) print(f"ATOM: {reply}") def main() -> None: parser = argparse.ArgumentParser(description="Chat with atom-native checkpoint") parser.add_argument("--checkpoint", required=True) parser.add_argument("--prompt", default="Bonjour") parser.add_argument("--max-packets", type=int, default=96) parser.add_argument("--max-length", type=int, default=256) parser.add_argument("--temperature", type=float, default=0.7) parser.add_argument("--top-k", type=int, default=8) parser.add_argument("--deterministic", action="store_true") parser.add_argument("--no-role-prime", action="store_true", help="do not wrap prompt as Utilisateur/Assistant") parser.add_argument( "--merge-ingest", action="store_true", help="enable MERGE during chat priming (default: off; train MERGE unchanged)", ) parser.add_argument("--device", default="cpu") parser.add_argument("--interactive", action="store_true") parser.add_argument( "--legacy-gen-labels", action="store_true", help="reproduce the pre-fix byte-tick decode (generated space/punct " "committed as whitespace/punctuation; docs/DECODE_FIX_AND_BPTT.md)", ) args = parser.parse_args() model = load_model(args.checkpoint, device=args.device) model.atomizer.legacy_gen_labels = bool(args.legacy_gen_labels) if args.interactive: interactive_loop(model, args) return reply = generate_reply( model, args.prompt, max_packets=args.max_packets, temperature=args.temperature, top_k=args.top_k, max_length=args.max_length, deterministic=args.deterministic, role_prime=not args.no_role_prime, merge_enabled=bool(args.merge_ingest), ) print(f"Prompt: {args.prompt}") print(f"Response: {reply}") print( f"[ingest] merge_enabled={bool(args.merge_ingest)} " f"atoms={len(model.core.atoms)} train_enable_merge={model.enable_merge}" ) if __name__ == "__main__": main()