""" Serves the trained model behind an OpenAI-shaped /v1/chat/completions endpoint so the existing harness can score it unchanged. ## Why a shim instead of a Python eval loop The scrambles, the scorer, the JSONL writer, the D1 upload path and the leaderboard are all in TypeScript, and the scramble seeds are index-addressed so every model faces the identical puzzles. Re-implementing scramble generation in Python to run a separate eval loop would risk silently diverging from those scrambles, which would quietly destroy the matched-pairs property that makes cross-model comparison meaningful. Pointing `callModel` at a local URL instead keeps one harness, one scorer and one code path, and the model becomes just another row in the results table. Nothing here is a general LLM server. It extracts the cube state from whatever prompt it is handed, greedily decodes a solution, and returns it in the response shape the harness already parses. """ import argparse, json, re, time import torch from transformers import LlamaForCausalLM import cube_tokenizer as T def extract_state(text: str) -> str: """Reads the 54-facelet state out of the benchmark prompt's net diagram. Parsed structurally rather than by regex over the whole prompt. The obvious shortcut -- strip everything that is not a facelet letter and take 54 of them -- silently produces a wrong state, because the surrounding prose is full of face letters on both sides of the diagram: the preamble names the colours ("U=White (Up), R=Red (Right)...") and the trailing instruction lists the move vocabulary ("(U, D, L, R, F, or B)"). Taking the first or last run picks up prose either way, and the result still looks like a valid 54-character state, so the failure would surface only as an unexplained low solve rate. Stickers are space-separated individually, so a net line is either 3 cells (a U or D row) or 12 cells (a row spanning L, F, R and B in that order). Facelet strings are ordered U R F D L B, so the wide rows are split by face and reassembled rather than concatenated in reading order. """ narrow, wide = [], [] for line in text.splitlines(): cells = line.split() if not cells or not all(len(c) == 1 and c in "URFDLB" for c in cells): continue if len(cells) == 3: narrow.append("".join(cells)) elif len(cells) == 12: wide.append("".join(cells)) if len(narrow) != 6 or len(wide) != 3: raise ValueError(f"unexpected net shape: {len(narrow)} narrow rows, {len(wide)} wide rows") up = "".join(narrow[:3]) down = "".join(narrow[3:]) left = "".join(r[0:3] for r in wide) front = "".join(r[3:6] for r in wide) right = "".join(r[6:9] for r in wide) back = "".join(r[9:12] for r in wide) return up + right + front + down + left + back @torch.no_grad() def _decode(model, state: str, device: str, temperature: float = 0.0, n: int = 1): """Decodes n solutions at once. temperature=0 is greedy (n is then pointless).""" prompt = torch.tensor([T.encode_state(state)] * n, device=device) ids = prompt finished = torch.zeros(n, dtype=torch.bool, device=device) for _ in range(T.MAX_SOLUTION + 1): logits = model(input_ids=ids).logits[:, -1, :] if temperature > 0: nxt = torch.multinomial(torch.softmax(logits / temperature, dim=-1), 1)[:, 0] else: nxt = logits.argmax(-1) nxt = torch.where(finished, torch.full_like(nxt, T.PAD), nxt) finished |= nxt == T.EOS ids = torch.cat([ids, nxt[:, None]], dim=1) if bool(finished.all()): break start = prompt.shape[1] return [T.decode_solution(row) for row in ids[:, start:].tolist()] def solve(model, state: str, device: str, attempts: int = 1, temperature: float = 1.5) -> str: """Returns a solution, verifying candidates against the engine when allowed more than one attempt. With attempts=1 this is a plain greedy decode -- the bare-model number, and the one to quote as the model's own solve rate. With attempts>1 the model proposes and the engine disposes: sampled candidates are each replayed, and the first that actually solves is returned. That is legitimate (the verifier is exact and free) but it is a *system* score, not a model score, so the two must be reported separately and never blended. It exists because a 20-move one-shot solve needs per-move accuracy around 0.9995 to hit 99% unaided, whereas a handful of verified samples buys the same result from a much weaker model. Temperature defaults to 1.5, which is high on purpose. Measured on the sanity checkpoint, 16 samples per state collapsed to only 2.8 distinct solutions at 0.7 and rescued *none* of the greedy failures; at 1.0, 4.0 distinct and 1 of 5; at 1.5, 11.0 distinct and 2 of 5. The model is confidently wrong rather than uncertain, so mild sampling just re-draws the same wrong answer -- the candidates have to be genuinely diverse before an exact verifier can do any work. A wrong candidate costs nothing here, since every one is checked. """ from gen_data import SOLVED, apply_sequence greedy = _decode(model, state, device)[0] if greedy and apply_sequence(state, greedy) == SOLVED: return " ".join(greedy) if attempts <= 1: return " ".join(greedy) for cand in _decode(model, state, device, temperature=temperature, n=attempts - 1): if cand and apply_sequence(state, cand) == SOLVED: return " ".join(cand) return " ".join(greedy) def solve_chunked(model, state: str, device: str, chunk: int = 8, max_rounds: int = 12, max_moves: int = 200): """Solves by repeatedly showing the model the *true* current state. The model proposes a full solution, only its first `chunk` moves are applied, and it is then shown the real resulting state and asked again. The returned value is the concatenation of every applied chunk -- a single move sequence from the original state, so it is scored exactly like a one-shot answer. ## Why this should work where more parameters and more layers did not A one-shot solve is blind after move 1: every later move reasons about a cube the model imagined, so a single early error makes the whole remainder correct reasoning about a state that does not exist. Success is roughly p^L for per-move accuracy p, which is why this model is ~100% at 3 moves and ~0% at 20 -- and why neither doubling parameters nor doubling depth moved that wall. Re-showing the true state resets the compounding: each round only needs the model to be right over `chunk` moves, which is the regime it is already good at. Predicting more than is executed is deliberate, and borrowed from robot action-chunking: the model plans a whole solution and only the confident front of it is used, so each applied move was chosen with the full plan in view rather than greedily. **This makes the task interactive rather than one-shot**, which is a different thing to measure. Results belong in their own mode and must never share a table with one-shot scores. """ from gen_data import SOLVED, apply_sequence applied: list[str] = [] current = state seen = {current} for _ in range(max_rounds): if current == SOLVED: break proposed = _decode(model, current, device)[0] if not proposed: break # Take the whole proposal when the engine confirms it finishes the job. # Truncating a correct answer to `chunk` is strictly harmful: it leaves the # cube one move short and forces another round, which is why chunk=8 turned # 9-move depth-6 solves into ~36-move ones at no gain in success rate. The # verifier is exact and free, so there is no reason to guess here. take = proposed if apply_sequence(current, proposed) == SOLVED else proposed[:chunk] current = apply_sequence(current, take) applied.extend(take) if len(applied) >= max_moves: break # No progress: the model is proposing something that returns to a state # already visited, so more rounds cannot help. if current in seen: break seen.add(current) return applied def build_app(model, device, model_name, args): from fastapi import FastAPI, Request app = FastAPI() @app.get("/v1/models") def models(): return {"object": "list", "data": [{"id": model_name, "object": "model"}]} @app.post("/v1/chat/completions") async def completions(req: Request): body = await req.json() text = "\n".join(m.get("content") or "" for m in body.get("messages", [])) t0 = time.time() try: answer = solve(model, extract_state(text), device, attempts=args.attempts) except ValueError as e: answer = f"ERROR: {e}" return { "id": "cube-local", "object": "chat.completion", "created": int(t0), "model": model_name, "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": answer}}], # Reported honestly: this model has no hidden reasoning, so the # harness should record zero rather than infer any from a blank field. "usage": {"prompt_tokens": T.MAX_SEQ, "completion_tokens": len(answer.split()), "total_tokens": T.MAX_SEQ + len(answer.split()), "completion_tokens_details": {"reasoning_tokens": 0}}, } return app def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument("--checkpoint", required=True, help="Local dir or Hub repo id.") p.add_argument("--port", type=int, default=8000) p.add_argument("--model-name", default="local/tiny-cube") p.add_argument("--attempts", type=int, default=1, help="1 = bare greedy model score. >1 samples extra candidates and " "keeps the first the engine confirms solves (a system score).") p.add_argument("--self-test", type=int, default=0, help="Solve N generated scrambles and exit, instead of serving.") args = p.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" model = LlamaForCausalLM.from_pretrained(args.checkpoint).to(device).eval() print(f"loaded {args.checkpoint} on {device}", flush=True) if args.self_test: import random from gen_data import SOLVED, apply_sequence, random_state rng = random.Random(0) solved = 0 for _ in range(args.self_test): state = random_state(rng, 200) moves = solve(model, state, device).split() if moves and apply_sequence(state, moves) == SOLVED: solved += 1 print(f"self-test: {solved}/{args.self_test} solved ({solved/args.self_test:.1%})") return import uvicorn uvicorn.run(build_app(model, device, args.model_name, args), host="0.0.0.0", port=args.port) if __name__ == "__main__": main()