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)


@torch.no_grad()
def solve_beam(model, state: str, device: str, beam: int = 16, top_k: int = 6):
    """Beam search over the model's move distribution, verified by the engine.

    Greedy decoding commits to the single most likely move at every step. Beam
    search keeps the `beam` highest-probability partial sequences instead, so a
    solution the model ranks second or third at one step is still reachable. The
    engine checks every completed candidate, so a wrong one costs nothing.

    This is the cheap half of the search hypothesis. Four attempts to fix the
    ~8-move wall by changing the *model* -- more parameters, more layers,
    interactive chunking, training on its own rollouts -- all failed, and what
    they had in common is the shape: one forward pass, one emitted sequence, no
    search. The literature does not use that shape. DeepCubeA learns a value
    function and runs A* over it; Kociemba is IDA* over pruning tables.

    Note what this does and does not add. It searches the model's *own*
    distribution, so it can recover a solution the model ranked highly but not
    first. It supplies no independent notion of distance-to-solved, so if the
    model's distribution is simply wrong at a given state -- as the rollout
    traces suggest it is beyond ~8 moves -- beam search has nothing better to
    find. A negative here therefore points at the value-function half of the
    hypothesis rather than closing it off.
    """
    from gen_data import SOLVED, apply_sequence

    prompt = T.encode_state(state)
    beams = [(0.0, [])]           # (cumulative logprob, move-token ids)
    finished: list[tuple[float, list[int]]] = []

    for _ in range(T.MAX_SOLUTION + 1):
        if not beams:
            break
        ids = torch.tensor([prompt + b[1] for b in beams], device=device)
        logprobs = torch.log_softmax(model(input_ids=ids).logits[:, -1, :], dim=-1)
        cand_lp, cand_ix = logprobs.topk(top_k, dim=-1)

        pool = []
        for (score, seq), lps, ixs in zip(beams, cand_lp.tolist(), cand_ix.tolist()):
            for lp, ix in zip(lps, ixs):
                if ix == T.EOS:
                    finished.append((score + lp, seq))
                else:
                    pool.append((score + lp, seq + [ix]))
        pool.sort(key=lambda x: -x[0])
        beams = pool[:beam]

    # Most probable first: the engine is exact, so the only cost of a wrong
    # candidate is the replay.
    for _, seq in sorted(finished, key=lambda x: -x[0]):
        moves = T.decode_solution(seq + [T.EOS])
        if moves and apply_sequence(state, moves) == SOLVED:
            return moves
    return T.decode_solution((finished[0][1] if finished else []) + [T.EOS])


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