"""ATOM chat Space: auto-loads the latest ATOM checkpoint from thefinalboss/atom-ai. ATOM is a byte-level model (persistent toroidal field + Atomizer byte tokenizer). No transformers / BPE. The model code (src/ + tools/chat_atom_native.py) is pulled from the model repo at startup so it stays in sync with the checkpoints. """ from __future__ import annotations import importlib.util import os import re import sys import threading import time import traceback from datetime import datetime, timezone import torch import gradio as gr from huggingface_hub import HfApi, hf_hub_download, snapshot_download MODEL_REPO = os.environ.get("ATOM_MODEL_REPO", "thefinalboss/atom-ai") REFRESH_SECONDS = int(os.environ.get("ATOM_REFRESH_SECONDS", str(30 * 60))) NUM_THREADS = int(os.environ.get("ATOM_NUM_THREADS", "2")) torch.set_num_threads(NUM_THREADS) api = HfApi() # --------------------------------------------------------------------------- # Model code from the model repo # --------------------------------------------------------------------------- def _load_code(): code_dir = os.environ.get("ATOM_CODE_DIR") if not code_dir: code_dir = snapshot_download( MODEL_REPO, allow_patterns=["src/*.py", "src/**/*.py", "tools/chat_atom_native.py"], ) if code_dir not in sys.path: sys.path.insert(0, code_dir) spec = importlib.util.spec_from_file_location( "chat_atom_native", os.path.join(code_dir, "tools", "chat_atom_native.py") ) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod chat_mod = _load_code() # --------------------------------------------------------------------------- # Latest-checkpoint picker # --------------------------------------------------------------------------- _STEP_RE = re.compile(r"step_(\d+)") def _step_of(path: str) -> int: m = _STEP_RE.search(os.path.basename(path)) return int(m.group(1)) if m else -1 def list_candidates(): """Return list of dicts {path, run, step, date, oid} for eligible .pt files.""" out = [] for it in api.list_repo_tree(MODEL_REPO, path_in_repo="checkpoints", recursive=True, expand=True): p = getattr(it, "path", "") if not p.endswith(".pt"): continue parts = p.split("/") if len(parts) < 3 or parts[1] == "RELEASE": continue # old pre-fix release checkpoints are excluded lc = getattr(it, "last_commit", None) date = lc.date if lc is not None and lc.date else datetime.fromtimestamp(0, timezone.utc) out.append({ "path": p, "run": parts[1], "step": _step_of(p), "date": date, "oid": getattr(it, "blob_id", None) or (lc.oid if lc else None), }) return out def pick_latest(cands=None): """Highest step wins. Skip RELEASE, smoke, and pair-probe runs. Date is only a tie-break so a newer small probe cannot hide the tip.""" cands = list_candidates() if cands is None else cands def skip(c): r = c["run"] return ( r == "RELEASE" or r.startswith("gpu_smoke") or "pairs" in r or "prompt" in r or "role" in r or "ans" in r or c["step"] < 0 ) pool = [c for c in cands if not skip(c)] if not pool: raise RuntimeError(f"no eligible .pt checkpoints in {MODEL_REPO}/checkpoints") pool.sort(key=lambda c: (c["step"], c["date"]), reverse=True) return pool[0] # --------------------------------------------------------------------------- # Model state (hot-swappable) # --------------------------------------------------------------------------- _lock = threading.Lock() # guards generation + swap (model is stateful) _state = {"model": None, "info": None, "loaded_at": None, "last_check": None, "error": None} def _load_checkpoint(cand): local = hf_hub_download(MODEL_REPO, cand["path"]) t0 = time.time() model = chat_mod.load_model(local, device="cpu") model.atomizer.legacy_gen_labels = False print(f"[atom-chat] loaded {cand['path']} step={cand['step']} in {time.time()-t0:.1f}s", flush=True) return model def refresh(force: bool = False) -> str: """Check the repo for a newer checkpoint and hot-swap if needed.""" try: cand = pick_latest() _state["last_check"] = datetime.now(timezone.utc) cur = _state["info"] if force or cur is None or cand["path"] != cur["path"] or cand["oid"] != cur["oid"]: model = _load_checkpoint(cand) # load outside the lock; swap atomically with _lock: _state["model"] = model _state["info"] = cand _state["loaded_at"] = datetime.now(timezone.utc) _state["error"] = None return status_md() except Exception as e: # keep serving the previous model traceback.print_exc() _state["error"] = f"{type(e).__name__}: {e}" return status_md() def status_md() -> str: info = _state["info"] if info is None: s = "**No checkpoint loaded yet.**" else: s = ( f"**Loaded checkpoint:** `{info['path']}` \n" f"**Run:** `{info['run']}` · **Step:** {info['step']:,} · " f"**Committed:** {info['date'].strftime('%Y-%m-%d %H:%M UTC')}" ) if _state["last_check"]: s += f" \n_Last checked for a newer checkpoint: {_state['last_check'].strftime('%Y-%m-%d %H:%M UTC')} (auto every {REFRESH_SECONDS//60} min)_" if _state["error"]: s += f" \n⚠️ Last refresh error: `{_state['error']}`" return s def _bg_loop(): while True: time.sleep(REFRESH_SECONDS) refresh() refresh() threading.Thread(target=_bg_loop, daemon=True).start() # --------------------------------------------------------------------------- # Chat # --------------------------------------------------------------------------- def _text_of(content) -> str: if isinstance(content, str): return content if isinstance(content, list): return " ".join(p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text") if isinstance(content, dict): return content.get("text", "") or "" return str(content or "") def _build_prompt(message: str, history, use_history: bool, max_hist_chars: int = 400) -> str: cur = f"Utilisateur: {message}\nAssistant:" if not use_history or not history: return cur lines = [] for turn in history: role = turn.get("role") if isinstance(turn, dict) else None text = _text_of(turn.get("content") if isinstance(turn, dict) else "").strip() if not text or role not in ("user", "assistant"): continue tag = "Utilisateur" if role == "user" else "Assistant" lines.append(f"{tag}: {text}") ctx = "" for line in reversed(lines): # keep most recent turns within the char budget if len(ctx) + len(line) + 1 > max_hist_chars: break ctx = line + "\n" + ctx return ctx + cur def respond(message, history, temperature, top_k, max_length, use_history, deterministic): message = _text_of(message).strip() if not message: return "(empty message)" prompt = _build_prompt(message, history, bool(use_history)) with _lock: model = _state["model"] if model is None: return "Model not loaded yet — please retry in a moment." # Model construction seeds torch's global RNG with a fixed value, which would make # every fresh process produce identical "random" samples; reseed per request. if not deterministic: torch.manual_seed(int.from_bytes(os.urandom(4), "little")) t0 = time.time() reply = chat_mod.generate_reply( model, prompt, max_packets=int(max_length), # byte-tick: 1 packet ~ 1 byte temperature=float(temperature), top_k=int(top_k), max_length=int(max_length), deterministic=bool(deterministic), role_prime=True, # prompt already carries Utilisateur/Assistant tags ) dt = time.time() - t0 # Trim a hallucinated next turn and a trailing partial UTF-8 sequence. for tag in ("\nUtilisateur:", "\nAssistant:", "\nUser:"): if tag in reply: reply = reply.split(tag, 1)[0] reply = reply.rstrip("\ufffd").rstrip() print(f"[atom-chat] reply in {dt:.2f}s ({len(reply)} chars)", flush=True) return reply if reply.strip() else "(ATOM produced only whitespace — try again or raise temperature)" DESCRIPTION = f""" **ATOM** is an experimental **byte-level** model: a persistent toroidal field plus the *Atomizer* byte tokenizer — no transformer, no BPE / subword tokenizer, ~480k parameters, running on a free CPU. - It is **bilingual French + English by design**, but the **current checkpoints are trained on French only** — prompts in French work best. - **Output is not coherent yet.** Expect fragments of French-looking text; this demo exists to watch the model improve as training progresses. - The Space automatically loads the **latest checkpoint** from the model repo [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO}) and re-checks for a newer one every {REFRESH_SECONDS//60} minutes. """ with gr.Blocks(title="ATOM chat") as demo: gr.Markdown("# ATOM — byte-level field model (experimental)") gr.Markdown(DESCRIPTION) with gr.Row(): status = gr.Markdown(status_md()) reload_btn = gr.Button("Reload latest checkpoint", scale=0) reload_btn.click(lambda: refresh(force=False), outputs=status, api_name="reload") demo.load(status_md, outputs=status, api_visibility="undocumented") gr.ChatInterface( respond, additional_inputs=[ gr.Slider(0.05, 2.0, value=0.7, step=0.05, label="Temperature"), gr.Slider(1, 64, value=8, step=1, label="Top-k"), gr.Slider(16, 512, value=256, step=16, label="Max length (bytes)"), gr.Checkbox(value=False, label="Include recent chat history in the prompt (short, ~400 chars)"), gr.Checkbox(value=False, label="Deterministic (greedy)"), ], additional_inputs_accordion=gr.Accordion("Generation settings", open=True), examples=[[m, 0.7, 8, 256, False, False] for m in ("Bonjour, comment ça va ?", "Qui es-tu ?", "Raconte-moi une histoire.")], cache_examples=False, api_name="chat", concurrency_limit=1, ) if __name__ == "__main__": demo.queue(default_concurrency_limit=1, max_size=32).launch()