system_one_games/runtime.py

"""Load one game adapter from the shared repository and score supplied options."""

from __future__ import annotations

import json
import math
from pathlib import Path, PurePosixPath
from typing import Any, Sequence

import torch

from . import _core as core

GAMES = ("dino", "snake", "tetris")


def _root(path: str | Path, game: str, revision: str | None) -> Path:
    if game not in GAMES:
        raise core.ModelError("game must be dino, snake or tetris")
    local = Path(path)
    if local.is_dir():
        return local
    from huggingface_hub import snapshot_download

    return Path(snapshot_download(str(path), revision=revision, allow_patterns=[
        "manifest.json", "checksums.sha256", f"{game}/*",
    ]))


class GameModel:
    """A fine-tuned Laya model. Inputs must include the game's engineered options."""

    def __init__(self, model: core.SystemOneDecisionModel, game: str):
        self.model, self.game = model, game

    @classmethod
    def from_pretrained(cls, path: str | Path = "getSTEAV/system-one-games", *, game: str,
                        revision: str | None = None, base: str | Path | None = None,
                        device: str | None = None, dtype: str = "auto") -> "GameModel":
        root = _root(path, game, revision)
        manifest = json.loads((root / "manifest.json").read_text())
        if manifest.get("schema") != "steav/system-one-games/v1":
            raise core.ModelError("unsupported game release manifest")
        for rel, expected in manifest["models"][game]["files"].items():
            relative = PurePosixPath(rel)
            target = root / relative
            # Hub snapshots use symlinks into their shared blob cache. Validate
            # the manifest path itself, then verify the referenced file's bytes.
            if relative.is_absolute() or ".." in relative.parts or not target.is_file():
                raise core.ModelError("missing or invalid release file: " + rel)
            if core._sha256(target) != expected:
                raise core.ModelError("release checksum mismatch: " + rel)
        model = core.SystemOneDecisionModel.from_pretrained(root / game, base=base,
                                                           device=device, dtype=dtype)
        return cls(model, game)

    def _temperature(self, question: dict) -> float:
        value = self.model.temperature(question)
        if not math.isfinite(value) or value <= 0:
            raise core.ModelError("calibration temperature must be finite and positive")
        return value

    @torch.inference_mode()
    def predict_many(self, rows: Sequence[dict], *, batch_size: int = 32) -> list[dict]:
        """Score rows {state, question}; each row may supply different options.

        Returns an option id and a calibrated probability per supplied option.
        Batch size 32 and original file order match the historical evaluation.
        """
        if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size < 1:
            raise core.ModelError("batch_size must be a positive integer")
        answers = []
        for start in range(0, len(rows), batch_size):
            chunk = [row.get("system_one", row) for row in rows[start:start + batch_size]]
            for row in chunk:
                if not isinstance(row.get("question"), dict) or row["question"].get("type") != "choice":
                    raise core.ModelError("game models require a choice question with supplied options")
            encoded = [self.model.encode(row["state"], row["question"]) for row in chunk]
            for row, item, logits in zip(chunk, encoded, self.model._logits(encoded)):
                values = torch.softmax(logits / self._temperature(row["question"]), -1).cpu().tolist()
                slots = core.question_slots(row["question"])
                if len(values) != len(slots) or not all(math.isfinite(p) and 0 <= p <= 1 for p in values):
                    raise core.ModelError("invalid decision probabilities")
                best = max(range(len(values)), key=values.__getitem__)
                answers.append({"choice": slots[best]["id"], "confidence": values[best],
                                "probabilities": {slot["id"]: p for slot, p in zip(slots, values)},
                                "tokens": len(item.ids), "dropped_state_tokens": item.dropped_state_tokens})
        return answers

    def predict(self, state: Any, question: dict) -> dict:
        return self.predict_many([{"state": state, "question": question}], batch_size=1)[0]


class StockLayaModel(GameModel):
    """The pinned unadapted Laya checkpoint, evaluated with the same input layout."""

    @classmethod
    def from_pretrained(cls, *, max_seq_length: int, base: str | Path | None = None,
                        device: str | None = None, dtype: str = "auto") -> "StockLayaModel":
        from safetensors.torch import load_file
        from transformers import AutoTokenizer

        base_dir = core._local_dir(base or core.LAYA_BASE_REPO, core.LAYA_BASE_REVISION,
                                  [*core.LAYA_BASE_FILES, "rl_agent_config.json"])
        agent = json.loads((base_dir / "rl_agent_config.json").read_text())
        device_obj = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        dtype_obj = core._resolve_dtype(dtype, device_obj)
        encoder = core.load_laya_encoder(base_dir, dtype_obj, device_obj, core.LAYA_BASE_SHA256)
        weights_path, _, tokenizer_dir = core.base_layout(base_dir)
        weights = {k: v for k, v in load_file(str(weights_path)).items() if not k.startswith("encoder.")}
        head = core.LayaDecisionHead(int(encoder.get_input_embeddings().weight.shape[1]),
                                    int(agent.get("head_layers", 2)),
                                    len(agent.get("act_costs", {"escalate": .5})) + 1).float()
        head.load_state_dict(weights, strict=True)
        head.to(device_obj).eval()
        tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_dir), trust_remote_code=False,
                                                  local_files_only=True)
        tokenizer.padding_side = "right"
        model = core.SystemOneDecisionModel(encoder.eval(), tokenizer, head,
                                            {"global_temperature": 1.0, "buckets": {}}, {}, {},
                                            device_obj, dtype_obj, max_seq_length, None, None, False)
        result = cls(model, "stock-laya")
        result.agent_config = agent
        return result

    def _temperature(self, question: dict) -> float:
        count = len(core.question_slots(question))
        size = "2" if count <= 2 else "3-5" if count <= 5 else "6-10" if count <= 10 else "11+"
        by_options = self.agent_config.get("temperature_by_options", {})
        by_type = self.agent_config.get("temperature", [1.0, 1.0, 1.0])
        value = float(by_options.get(question["type"] + ":" + size,
                                     by_type[core.LAYA_QTYPES[question["type"]]]))
        if not math.isfinite(value) or value <= 0:
            raise core.ModelError("stock calibration temperature must be finite and positive")
        return value