File size: 7,238 Bytes
9220532 | 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 | """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
|