code/serve.py

"""
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 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()