| |
| """Interactive Terminal Chat for `o1` (3.32M params, 32,512 B O(1) carrier + online GSSM).""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
|
|
| HERE = Path(__file__).resolve().parent |
| if str(HERE) not in sys.path: |
| sys.path.insert(0, str(HERE)) |
|
|
| from o1_native.organism import O1Organism |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description="Interactive Terminal Chat with o1") |
| ap.add_argument("--checkpoint", type=str, default=str(HERE / "o1.pt")) |
| ap.add_argument("--device", type=str, default="cpu", help="cpu is fastest for B=1 chat; cuda/mps also supported") |
| ap.add_argument("--no-learn", action="store_true", help="Disable live per-turn online weight learning") |
| ap.add_argument("--max-tokens", type=int, default=96) |
| ap.add_argument("--temperature", type=float, default=0.0) |
| args = ap.parse_args() |
|
|
| ckpt = Path(args.checkpoint) |
| payload = torch.load(ckpt, map_location="cpu", weights_only=False) |
| meta = payload.get("meta", {}) |
| org = O1Organism.load_snapshot(ckpt, device_override=args.device) |
| learn_mode = not args.no_learn |
|
|
| n_params = sum(p.numel() for p in org.model.parameters()) |
| print(f"o1 | params={n_params:,} | carrier={org.carrier.nbytes():,} B | learn={'on' if learn_mode else 'off'} | dev={args.device}") |
| if meta: |
| print(f"by {meta.get('author', 'David Tom Foss')} | {meta.get('email', 'david@wirku.ng')} | {meta.get('x', 'https://x.com/fossdt')}") |
| print("Commands: /reset /learn on|off /stats /save <path> /exit\n") |
|
|
| while True: |
| try: |
| msg = input("you> ").strip() |
| except (EOFError, KeyboardInterrupt): |
| print() |
| break |
| if not msg: |
| continue |
| if msg in ("/exit", "/quit"): |
| break |
| if msg == "/reset": |
| org.reset_carrier() |
| print("[o1] carrier reset to zeros (32,512 B)") |
| continue |
| if msg == "/learn on": |
| learn_mode = True |
| print("[o1] live online learning: ON") |
| continue |
| if msg == "/learn off": |
| learn_mode = False |
| print("[o1] live online learning: OFF") |
| continue |
| if msg == "/stats": |
| wb = org.causal_store.verify_weltbuch() |
| print( |
| f"[o1] tokens_seen={org.total_tokens_seen:,} | plastic_updates={org.plastic_updates} " |
| f"| carrier_bytes={org.carrier.nbytes():,} | 16d_drift={org.invariant_16d_drift():.2e} " |
| f"| weltbuch_valid={wb.get('valid')} ({wb.get('total_triples')} triples)" |
| ) |
| continue |
| if msg.startswith("/save "): |
| target = msg[6:].strip() |
| if target: |
| org.save_snapshot(target) |
| print(f"[o1] saved state -> {target}") |
| continue |
|
|
| turn_prompt = f"User: {msg}\nAssistant:" |
| t0 = time.perf_counter() |
| out = org.generate( |
| turn_prompt, |
| max_new_tokens=args.max_tokens, |
| temperature=args.temperature, |
| use_live_carrier=True, |
| stop_on_newline=False, |
| ) |
| ans = str(out.get("answer", "")).split("\nUser:", 1)[0].strip() |
| dt = max(time.perf_counter() - t0, 1e-6) |
| toks = max(int(out.get("emitted_token_count", 0)), 1) |
|
|
| step_rep = org.stream_step(turn_prompt, learn=learn_mode) |
| _ = org.stream_step(f" {ans}\n", learn=False) |
| route = out.get("organ_route") or "TEXT" |
| loss_nats = float(step_rep.get("metrics", {}).get("loss_nats", 0.0)) |
| print(f"o1 [{route} | {toks/dt:.1f} tok/s | loss={loss_nats:.3f}]> {ans}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
