File size: 7,245 Bytes
e29f88c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """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()
|