scripts/demo/inference.py

"""Stream native loop observations without reference-state feedback."""

from dataclasses import dataclass
import math
from queue import Queue
from threading import Event, Thread
from time import perf_counter
from typing import Any, Iterator

import torch

from scripts.dataset.pointer import SYMBOLS, execute, render_prompt


SAFETY_CAP = 272
STOP_THRESHOLD = 0.5


@dataclass(frozen=True)
class DemoRequest:
    """Ordered full A–Z table; order is part of the model's actual prompt."""

    pairs: tuple[tuple[str, str], ...]
    start: str
    steps: int

    @classmethod
    def parse(cls, value: Any, *, max_steps: int = 256) -> "DemoRequest":
        if not isinstance(value, dict) or set(value) != {"mapping", "start", "steps"}:
            raise ValueError("Supply exactly mapping, start, and steps")
        pairs = value["mapping"]
        if not isinstance(pairs, list) or len(pairs) != 26:
            raise ValueError("Supply all 26 A–Z rules, with one destination per node")
        if any(not isinstance(p, list) or len(p) != 2 or
               any(not isinstance(s, str) or s not in SYMBOLS for s in p) for p in pairs):
            raise ValueError("Each rule must be a pair of uppercase A–Z letters")
        if {p[0] for p in pairs} != set(SYMBOLS):
            raise ValueError("Each A–Z source must occur exactly once")
        if not isinstance(value["start"], str) or value["start"] not in SYMBOLS:
            raise ValueError("Start must be an uppercase A–Z letter")
        if type(value["steps"]) is not int or not 1 <= value["steps"] <= max_steps:
            raise ValueError(f"Steps must be an integer from 1 to {max_steps}")
        return cls(tuple(tuple(p) for p in pairs), value["start"], value["steps"])

    def payload(self) -> dict:
        return {"mapping": [list(p) for p in self.pairs], "start": self.start, "steps": self.steps}

    def prompt(self) -> str:
        return render_prompt([list(p) for p in self.pairs], self.start, self.steps)


class InferenceCancelled(Exception):
    """The client closed its generator; stop at the next completed readout."""


def score_trace(request: DemoRequest, rows: list[dict], stop_reason: str) -> dict:
    """Final letter, strict requested prefix and timing are separate metrics."""
    expected = execute(dict(request.pairs), request.start, request.steps)
    predictions = [row["prediction"] for row in rows]
    exact_stop = stop_reason == "head" and len(rows) == request.steps
    correct = bool(predictions) and predictions[-1] == expected[-1]
    complete = len(predictions) >= request.steps and predictions[:request.steps] == expected
    return {"answer_correct": correct, "complete_trajectory": complete,
            "exact_stop": exact_stop, "joint_success": correct and exact_stop,
            "stop_reason": stop_reason, "executed_loops": len(rows),
            "first_error": next((r["loop"] for r in rows[:request.steps] if not r["correct"]), None)}


def stream_inference(model: Any, tokenizer: Any, spec: dict, request: DemoRequest,
                     *, checkpoint_id: str, cancel: Event | None = None,
                     loop_cap: int | None = None) -> Iterator[dict]:
    """Yield cumulative JSON snapshots as loops arrive, then a final summary.

    One worker owns one native forward. The bounded 272-loop cap also bounds
    the queue. A consumer must serialize calls to its model (Gradio concurrency
    one). Closing the iterator cancels and joins the worker before releasing GPU
    ownership. Targets are computed only in the consumer, never passed to R.
    """
    if model.training or model.completion_head is None:
        raise ValueError("Demo requires an eval-mode model with learned completion")
    loop_cap = SAFETY_CAP if loop_cap is None else loop_cap
    if type(loop_cap) is not int or loop_cap < 1:
        raise ValueError("Loop cap must be a positive integer")
    cancel = cancel if cancel is not None else Event()
    queue: Queue = Queue()
    rows: list[dict] = []
    mapping = dict(request.pairs)
    reference = request.start
    previous_prediction = request.start
    started = perf_counter()
    metadata = {"schema_version": 1, "source": "live", "checkpoint": checkpoint_id,
                "request": request.payload(), "threshold": STOP_THRESHOLD, "safety_cap": loop_cap}

    def observe(loop, scores, stop_logits, state_logits):
        if cancel.is_set():
            raise InferenceCancelled()
        scores = scores[0].float().cpu()
        probability = float(stop_logits.sigmoid().item())
        if not torch.isfinite(scores).all() or not math.isfinite(probability):
            raise FloatingPointError("Nonfinite model output")
        top = scores.topk(3)
        probabilities = scores.softmax(-1)
        queue.put(("loop", {"loop": loop, "prediction": SYMBOLS[int(scores.argmax())],
                   "direct_prediction": SYMBOLS[int(state_logits[0].argmax())] if state_logits is not None else None,
                   "stop_probability": probability,
                   "top3": [{"symbol": SYMBOLS[i], "probability": float(probabilities[i])}
                            for i in top.indices.tolist()],
                   "arrival_seconds": perf_counter() - started}))

    def run():
        try:
            if cancel.is_set():
                raise InferenceCancelled()
            device = next(model.parameters()).device
            encoded = tokenizer(request.prompt(), add_special_tokens=False, return_tensors="pt")
            with torch.inference_mode():
                result = model(input_ids=encoded["input_ids"].to(device),
                               attention_mask=encoded["attention_mask"].to(device),
                               num_loops=loop_cap, completion_threshold=STOP_THRESHOLD,
                               readout_token_ids=spec["token_ids"], loop_callback=observe)
            signaled = result.stop_logits[0, -1].sigmoid().item() >= STOP_THRESHOLD
            queue.put(("done", "head" if signaled else "cap"))
        except InferenceCancelled:
            queue.put(("cancelled", None))
        except Exception as error:
            # Transport the original failure to the owning generator, never hide it.
            queue.put(("error", error))

    worker = Thread(target=run, name="loopformer-inference", daemon=True)
    worker.start()
    try:
        while True:
            kind, value = queue.get()
            if kind == "error":
                raise value
            if kind == "cancelled":
                return
            if kind == "done":
                yield {**metadata, "status": "complete", "trace": list(rows),
                       "summary": score_trace(request, rows, value),
                       "inference_seconds": perf_counter() - started}
                return
            reference = mapping[reference]
            value.update(expected=reference, correct=value["prediction"] == reference,
                         edge_valid=mapping[previous_prediction] == value["prediction"])
            previous_prediction = value["prediction"]
            rows.append(value)
            yield {**metadata, "status": "streaming", "trace": list(rows)}
    finally:
        cancel.set()
        worker.join()