"""Local inference server: serve a trained GPCT checkpoint to the harness. Milestone 6's "small local inference server behind predict(ctx)". Downloads ckpt.pt + meta.pkl from your HF model repo, loads them with the SAME pinned nanoGPT the training job used, and exposes the model over localhost HTTP. The browser harness auto-detects it via js/remote-model.js (mock model stays the fallback when the server is down). pip install torch huggingface_hub python tools/model_server.py # serves shazmate/gpct-trial python tools/model_server.py --repo you/chess-gpt --model-elo 1600 python tools/model_server.py --local ./out # use a local ckpt dir instead Then open the harness (node tools/serve.mjs) — the page picks the server up on load; the console logs which model is live. Restart this server after a new checkpoint is pushed to pick it up (hf_hub_download re-checks the remote). Context contract (must mirror engine.js + nerf_batch.py — the bare-history rule): the model context is ...historyTokens [nerf token if quality] ...moveTokens historyTokens arrive BARE from the harness (never contain nerf tokens); the current move's own nerf conditioning arrives as `quality`. The Elo slot for the side the model plays comes from --model-elo ('any' -> , or a number -> its 200-point bucket); the human's slot from the rating typed into the setup form (missing -> ). API: GET /health -> {repo, iter_num, best_val_loss, params, model_args} POST /predict -> body {historyTokens, quality, moveTokens, turn, opponentRating}; returns a JSON array of vocab_size probabilities for the next token. """ from __future__ import annotations import argparse import importlib.util import json import os import pickle import threading import time import urllib.request from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer # Same commit train_chess_hf.py trains with — architecture code must match. NANO_COMMIT = "3adf61e154c3fe3fca428ad6bc3818b27a3b8291" CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "chess-gpt") QUALITY_TO_NERF = {"inaccuracy": "", "mistake": "", "blunder": ""} def elo_bucket_token(elo) -> str: """Numeric Elo -> bucket token; mirrors chess-tokeniser/vocab.py.""" try: e = int(elo) except (TypeError, ValueError): return "" if e < 800: return "" if e >= 3000: return "" return f"" def fetch_nanogpt_model_module(): """Import the pinned nanoGPT model.py (cached under ~/.cache/chess-gpt).""" os.makedirs(CACHE_DIR, exist_ok=True) path = os.path.join(CACHE_DIR, f"nanogpt_model_{NANO_COMMIT[:8]}.py") if not os.path.exists(path): url = f"https://raw.githubusercontent.com/karpathy/nanoGPT/{NANO_COMMIT}/model.py" print(f"== fetching nanoGPT model.py @ {NANO_COMMIT[:8]}") with open(path, "wb") as f: f.write(urllib.request.urlopen(url).read()) spec = importlib.util.spec_from_file_location("nanogpt_model", path) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod def load_model(args): import torch if args.local: ckpt_path = os.path.join(args.local, "ckpt.pt") meta_path = os.path.join(args.local, "meta.pkl") else: from huggingface_hub import hf_hub_download from huggingface_hub.errors import EntryNotFoundError print(f"== downloading checkpoint from {args.repo}") ckpt_path = hf_hub_download(args.repo, "ckpt.pt") try: meta_path = hf_hub_download(args.repo, "meta.pkl") except EntryNotFoundError: # Mid-run periodic pushes carry only ckpt.pt; meta.pkl lands with # the final upload. The dataset repo holds the identical file. print(f"== meta.pkl not in {args.repo} yet — using {args.meta_repo}") meta_path = hf_hub_download(args.meta_repo, "meta.pkl", repo_type="dataset") with open(meta_path, "rb") as f: meta = pickle.load(f) nano = fetch_nanogpt_model_module() try: ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True) except Exception: ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) # our own file model_args = ckpt["model_args"] model = nano.GPT(nano.GPTConfig(**model_args)) state = {k.removeprefix("_orig_mod."): v for k, v in ckpt["model"].items()} model.load_state_dict(state) model.eval() info = { "repo": args.local or args.repo, "iter_num": int(ckpt.get("iter_num", -1)), "best_val_loss": float(ckpt.get("best_val_loss", float("nan"))), "params": sum(p.numel() for p in model.parameters()), "model_args": {k: v for k, v in model_args.items()}, "vocab_size": meta["vocab_size"], } print(f"== model ready: {info['params']:,} params, iter {info['iter_num']}, " f"best val loss {info['best_val_loss']:.4f}") return model, meta, info def main(): ap = argparse.ArgumentParser() ap.add_argument("--repo", default="shazmate/gpct-trial", help="HF model repo with ckpt.pt + meta.pkl") ap.add_argument("--local", default=None, help="local dir with ckpt.pt + meta.pkl (skips download)") ap.add_argument("--meta-repo", default="shazmate/lichess-chess-tokens", help="dataset repo to fetch meta.pkl from when the model repo lacks it (mid-run pushes)") ap.add_argument("--port", type=int, default=8123) ap.add_argument("--model-elo", default="any", help="Elo the model plays at: a number for its bucket, 'any' for ") args = ap.parse_args() import torch model, meta, info = load_model(args) stoi, vocab_size = meta["stoi"], meta["vocab_size"] block_size = info["model_args"].get("block_size", 512) model_elo_token = "" if args.model_elo == "any" else elo_bucket_token(args.model_elo) lock = threading.Lock() def build_ids(req): history = req.get("historyTokens") or [] move_tokens = req.get("moveTokens") or [] quality = req.get("quality") turn = req.get("turn", "w") human_elo_token = elo_bucket_token(req.get("opponentRating")) white = model_elo_token if turn == "w" else human_elo_token black = model_elo_token if turn == "b" else human_elo_token names = ["", white, black, *history] if quality: names.append(QUALITY_TO_NERF[quality]) names.extend(move_tokens) unknown = [t for t in names if t not in stoi] if unknown: raise ValueError(f"tokens not in vocab: {unknown}") return [stoi[t] for t in names] class Handler(BaseHTTPRequestHandler): def log_message(self, *a): # quiet the default per-request stderr noise pass def _send(self, code, payload): body = json.dumps(payload).encode() self.send_response(code) self.send_header("Content-Type", "application/json") self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Headers", "content-type") self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) def do_OPTIONS(self): self.send_response(204) self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") self.send_header("Access-Control-Allow-Headers", "content-type") self.end_headers() def do_GET(self): if self.path == "/health": self._send(200, info) else: self._send(404, {"error": "unknown path"}) def do_POST(self): if self.path != "/predict": return self._send(404, {"error": "unknown path"}) try: req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)))) ids = build_ids(req)[-block_size:] t0 = time.time() with lock, torch.no_grad(): idx = torch.tensor([ids], dtype=torch.long) logits, _ = model(idx) # (1, 1, vocab) at the last position probs = torch.softmax(logits[0, -1, :vocab_size], dim=-1) print(f"predict: {len(ids)} ctx tokens -> {1000 * (time.time() - t0):.0f} ms") self._send(200, probs.tolist()) except Exception as e: # noqa: BLE001 — surface the reason to the harness self._send(400, {"error": str(e)}) # 0.0.0.0: this copy runs inside a Hugging Face Space container, which # must accept traffic from the Space's reverse proxy, not just localhost. server = ThreadingHTTPServer(("0.0.0.0", args.port), Handler) print(f"== serving on http://0.0.0.0:{args.port} (health: /health)") server.serve_forever() if __name__ == "__main__": main()