Spaces:
Running
Running
| """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() | |