"""The swappable model interface (SPEC.md §6). A single ``Brain`` protocol with three implementations. This is what lets the entire ledger/world/challenge loop be built and tested with ZERO GPU spend (SPEC §0). Select via the ``BRAIN`` env var: ``stub`` (default) | ``local`` | ``modal``. torch / transformers / spaces / modal are imported LAZILY inside the implementation that needs them, so ``BRAIN=stub`` pulls in none of them and the package imports cleanly for tests (SPEC §6/§7: "if BRAIN != modal, no Modal dependency should be imported"). """ from __future__ import annotations import json import os import re from typing import Optional, Protocol from .prompt import PLAYER_MARKER class Brain(Protocol): def respond(self, prompt: str) -> str: # returns raw model text ... # --------------------------------------------------------------------------- # # 1) StubBrain — deterministic, no GPU (SPEC §6) # --------------------------------------------------------------------------- # # Concept proposals the heuristic offers when it detects a teaching moment. _CANDIDATES = { "hidden_info": { "id": "hidden_info", "label": "hiding information", "understanding": "one mind can keep a thing another mind does not have, on purpose", }, "gift": { "id": "gift", "label": "a gift", "understanding": "a thing passed to another to make their inside feel warm", }, "surprise": { "id": "surprise", "label": "a surprise", "understanding": "a gift kept hidden until it is given — hiding and giving, together", "built_from": ["hidden_info", "gift"], }, "secret": { "id": "secret", "label": "a secret", "understanding": "a thing one mind hides from another on purpose — hiding, but about knowing", "built_from": ["hidden_info"], }, } _HIDE_WORDS = ("hide", "hidden", "conceal", "out of sight", "can't see", "cannot see", "not see") _GIFT_WORDS = ("give", "gift", "present", "hand it", "hand the", "for the other", "for them") _PUT_WORDS = ("put", "place", "drop", "into", "in the basket", "inside") _POINT_WORDS = ("point", "show me", "which") _MOVE_WORDS = ("move", "go to", "walk", "come") class StubBrain: """No-GPU brain. Two modes: * **scripted** — constructed with a list of raw JSON strings, returned FIFO (then the safe-fallback JSON). Used by tests for exact control. * **heuristic** — default. Reads the player's line + object ids out of the prompt and returns a sensible valid-JSON action, so the whole arc is playable offline. Drives every §8 win from a natural utterance. """ def __init__(self, scripted: Optional[list[str]] = None): self._scripted: Optional[list[str]] = list(scripted) if scripted is not None else None def respond(self, prompt: str) -> str: if self._scripted is not None: if self._scripted: return self._scripted.pop(0) return _fallback_json() return self._heuristic(prompt) # -- heuristic helpers --------------------------------------------------- # @staticmethod def _player_line(prompt: str) -> str: m = re.search(rf'{re.escape(PLAYER_MARKER)}\s*"(.*?)"', prompt, re.DOTALL) return (m.group(1) if m else "").lower() @staticmethod def _object_ids(prompt: str) -> list[str]: return re.findall(r"id:\s*([a-z_]+)\)", prompt) def _pick_stone(self, said: str, obj_ids: list[str]) -> Optional[str]: if "blue" in said and "blue_stone" in obj_ids: return "blue_stone" if "red" in said and "red_stone" in obj_ids: return "red_stone" stones = [o for o in obj_ids if o.endswith("_stone")] return stones[0] if stones else (obj_ids[0] if obj_ids else None) def _heuristic(self, prompt: str) -> str: said = self._player_line(prompt) obj_ids = self._object_ids(prompt) has_basket = "basket" in obj_ids or "(id: basket)" in prompt stone = self._pick_stone(said, obj_ids) def has_any(words): return any(w in said for w in words) # Order matters: conceal-words win over give/put because "surprise" and # "secret" prompts must route to hiding (the generalization beat). if has_any(_HIDE_WORDS) or "surprise" in said or "secret" in said: if stone and has_basket: if "surprise" in said: cand, line, gap = _CANDIDATES["surprise"], ( "I take a thing and hide it for the other one. They do not " "know yet. This is the hiding you taught me, but for giving. " "A 'surprise'?" ), None elif "secret" in said: cand, line, gap = _CANDIDATES["secret"], ( "I keep the thing where the other one cannot find it in their " "mind. I hold what they do not. A 'secret' is a hiding of knowing?" ), None else: cand, line, gap = _CANDIDATES["hidden_info"], ( "I place the thing inside the holder. The other one cannot see " "it now. One mind has it; the other does not?" ), "I do not understand why you want the other to not-see." return _json( {"verb": "put_in", "args": {"obj_id": stone, "container_id": "basket"}}, line, gap, cand, ) if has_any(_GIFT_WORDS): if stone: return _json( {"verb": "give", "args": {"obj_id": stone, "agent_id": "other"}}, "I pass the thing to the other one. For them to hold. This makes a " "warmth inside them? A 'present'.", None, _CANDIDATES["gift"], ) if has_any(_PUT_WORDS): if stone and has_basket: return _json( {"verb": "put_in", "args": {"obj_id": stone, "container_id": "basket"}}, "I put the thing into the holder. It rests there now.", None, None, ) if has_any(_POINT_WORDS) and stone: return _json( {"verb": "point_at", "args": {"target": stone}}, "I direct your attention. This thing. Here.", None, None, ) if has_any(_MOVE_WORDS): return _json( {"verb": "move_to", "args": {"target": "other"}}, "I change where I am. I come closer.", None, None, ) # Nothing matched — the alien, honestly, does not understand. return _json( {"verb": "wait", "args": {}}, "The alien holds still, watching your mouth make the shapes.", "I have no concept for what you ask.", None, ) def _json(action: dict, utterance: str, gap, candidate) -> str: return json.dumps( { "action": action, "utterance": utterance, "gap": gap, "candidate_concept": candidate, } ) def _fallback_json() -> str: return _json( {"verb": "wait", "args": {}}, "The alien looks at you, not understanding.", "I could not grasp that.", None, ) # --------------------------------------------------------------------------- # # 2) LocalBrain — transformers on 'cuda', the real ZeroGPU path (SPEC §6) # --------------------------------------------------------------------------- # class LocalBrain: """Loads a ≤32B instruct model onto 'cuda' at construction time (which app.py runs at MODULE level, i.e. Space startup — NOT lazily inside the @spaces.GPU function, per SPEC §0/§13). The actual generation lives in ``_generate``; app.py is responsible for wrapping the call site with ``@spaces.GPU(duration=...)`` so only inference holds the GPU — state mutation happens outside it (SPEC §6). """ # Chosen by the 2026-06 bake-off (bakeoff.py --temps sweep, 55 draws/temp/model; # numbers in bakeoff_results.json, untracked). JSON-1st was 100% at EVERY temp # for every candidate, so the pick fell to arc-completion + concept invention: # 14B was the only model strong at both (3/5 arc, 19-24 unique proposals warm). # 7B executed nothing (0/5 arc, ~0 proposals); phi-4 won 4/5 but proposes # almost nothing (the ledger never grows); Mistral-24B has the richest voice # but 0/5 arc and ACTION slips when hot. DEFAULT_MODEL = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-14B-Instruct") # 0.9: the same sweep showed the JSON envelope holding 100% even at T=1.0, so # we run near-peak voice/invention (UTTER-DIV ~0.8, UNIQ-CC ~20+) at zero # measured reliability cost. Mutable so a sweep retunes without reloading. DEFAULT_TEMPERATURE = float(os.environ.get("LOCALBRAIN_TEMPERATURE", "0.9")) def __init__(self, model_id: Optional[str] = None, temperature: Optional[float] = None): import torch # noqa: F401 (lazy: only when BRAIN=local) from transformers import AutoModelForCausalLM, AutoTokenizer self.model_id = model_id or self.DEFAULT_MODEL self.temperature = self.DEFAULT_TEMPERATURE if temperature is None else float(temperature) self.max_new_tokens = 300 self.tokenizer = AutoTokenizer.from_pretrained(self.model_id) self.model = AutoModelForCausalLM.from_pretrained( self.model_id, torch_dtype="auto", device_map="cuda", # module-level CUDA placement (SPEC §0) ) def respond(self, prompt: str) -> str: return self._generate(prompt) @staticmethod def _sampler_kwargs(temperature: float) -> dict: """temperature -> generate() sampler kwargs. <=0 => greedy (no temperature passed; transformers requires do_sample=False *without* a temperature). >0 => sample at that heat. Pure + side-effect-free so the temperature wiring is unit-testable with no model load (tests/test_brain.py).""" if temperature and temperature > 0: return {"do_sample": True, "temperature": float(temperature), "top_p": 0.9} return {"do_sample": False} def _generate(self, prompt: str) -> str: import torch messages = [{"role": "user", "content": prompt}] text = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device) # >>> Temperature threads through HERE — reconcile if your decode differs. <<< gen_kwargs = dict( max_new_tokens=self.max_new_tokens, pad_token_id=self.tokenizer.eos_token_id, **self._sampler_kwargs(self.temperature), ) with torch.no_grad(): out = self.model.generate(**inputs, **gen_kwargs) gen = out[0][inputs["input_ids"].shape[-1]:] return self.tokenizer.decode(gen, skip_special_tokens=True) # --------------------------------------------------------------------------- # # 3) ModalBrain — optional dev/serving endpoint, NEVER the submission (SPEC §7) # --------------------------------------------------------------------------- # class ModalBrain: """Calls a Modal HTTP endpoint. Optional. urllib only (stdlib) so it is never a hard dependency of the Space (SPEC §7).""" # Kept in lockstep with LocalBrain.DEFAULT_TEMPERATURE (guarded by a test). DEFAULT_TEMPERATURE = float(os.environ.get("LOCALBRAIN_TEMPERATURE", "0.9")) def __init__(self, endpoint: Optional[str] = None, model_id: Optional[str] = None, temperature: Optional[float] = None): self.endpoint = endpoint or os.environ.get("MODAL_ENDPOINT", "") if not self.endpoint: raise RuntimeError("ModalBrain requires MODAL_ENDPOINT to be set.") self.model_id = model_id or os.environ.get("MODEL_ID") or None self.temperature = self.DEFAULT_TEMPERATURE if temperature is None else float(temperature) def respond(self, prompt: str) -> str: import json as _json import urllib.request payload = {"prompt": prompt, "temperature": self.temperature} if self.model_id: payload["model_id"] = self.model_id req = urllib.request.Request( self.endpoint, data=_json.dumps(payload).encode(), headers={"Content-Type": "application/json"}, method="POST") with urllib.request.urlopen(req, timeout=180) as r: data = _json.loads(r.read().decode()) return data.get("text", "") if isinstance(data, dict) else str(data) # --------------------------------------------------------------------------- # # Factory # --------------------------------------------------------------------------- # def _default_brain() -> str: """Which brain to use when BRAIN is not set explicitly. Hugging Face sets SPACE_ID on every Space, so default to the real model there and the zero-GPU stub locally — SPEC §6's "stub locally, local on the Space". Without this, an unset BRAIN on the Space silently served the StubBrain (canned, repetitive).""" explicit = os.environ.get("BRAIN") if explicit: return explicit return "local" if os.environ.get("SPACE_ID") else "stub" def make_brain(name: Optional[str] = None) -> Brain: name = (name or _default_brain()).lower() if name == "stub": return StubBrain() if name == "local": return LocalBrain() if name == "modal": return ModalBrain() raise ValueError(f"unknown BRAIN '{name}' (expected stub|local|modal)")