""" 1000 Rooms — Escape-Room game backend Served via gradio.Server (FastAPI + Gradio queuing + ZeroGPU support). """ import os import time import re import json import logging from pathlib import Path from huggingface_hub import hf_hub_download import spaces from gradio import Server from fastapi.responses import HTMLResponse log = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) # ─── Constants ──────────────────────────────────────────────────────────────── _NO_THINK_PREFIX = "/nothink\n\n" _GGUF_REPO = "build-small-hackathon/Nemotron-nano-4b-escape-room" _GGUF_FILENAME = "nemotron-room-lora-Q4_K_M.gguf" _ROOM_REQUIRED = { "room_name", "room_story", "room_prompt", "door_description", "door_prompt", "door_key_name", "door_key_prompt", "containers", "keys", } # ─── Prompts ────────────────────────────────────────────────────────────────── GENERATE_GAME_PROMPT = r"""You are a creative dungeon-master AI that generates escape-room content. Always respond with a single valid JSON object and nothing else — no markdown, no explanation, no preamble. Complete every placeholder: {"room_name":"", "room_story":"", "room_prompt":"", "door_description":"", "door_prompt":"", "door_key_name":"", "door_key_prompt":"", "containers":[ {"container_name":"","container_prompt":""}, {"container_name":"","container_prompt":""}, {"container_name":"","container_prompt":""}, {"container_name":"","container_prompt":""} ], "keys":[ {"key_name":"","key_prompt":""}, {"key_name":"","key_prompt":""} ]}""" CONTINUE_GAME_PROMPT = r"""You are a creative dungeon-master AI. The player has found the correct key and opened a container. Narrate the discovery vividly — describe opening the container and reveal item_to_give. Respond with a single valid JSON object and nothing else: {"text":""}""" OPEN_DOOR_PROMPT = r"""You are a creative dungeon-master AI. The player has used the correct key to open the exit door. Describe their triumphant escape in vivid detail. Respond with a single valid JSON object and nothing else: {"text":""}""" # ─── JSON schema for constrained sampling ───────────────────────────────────── _ROOM_SCHEMA = { "type": "object", "properties": { "room_name": {"type": "string"}, "room_story": {"type": "string"}, "room_prompt": {"type": "string"}, "door_description": {"type": "string"}, "door_prompt": {"type": "string"}, "door_key_name": {"type": "string"}, "door_key_prompt": {"type": "string"}, "containers": { "type": "array", "items": { "type": "object", "properties": { "container_name": {"type": "string"}, "container_prompt": {"type": "string"}, }, "required": ["container_name", "container_prompt"], }, "minItems": 4, }, "keys": { "type": "array", "items": { "type": "object", "properties": { "key_name": {"type": "string"}, "key_prompt": {"type": "string"}, }, "required": ["key_name", "key_prompt"], }, "minItems": 2, }, }, "required": list(_ROOM_REQUIRED), } _TEXT_SCHEMA = { "type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"], } # ─── Model loading ───────────────────────────────────────────────────────────── _model_path: str | None = None def _ensure_gguf() -> str: global _model_path if _model_path and Path(_model_path).exists(): return _model_path log.info("Downloading GGUF from HF Hub…") _model_path = hf_hub_download(repo_id=_GGUF_REPO, filename=_GGUF_FILENAME) log.info("GGUF cached at %s", _model_path) return _model_path def load_text_model(): from llama_cpp import Llama from llama_cpp.llama_chat_format import Jinja2ChatFormatter log.info("load_text_model: loading onto GPU…") with open("nemotron-template.jinja") as f: template_str = f.read() formatter = Jinja2ChatFormatter( template=template_str, eos_token="<|im_end|>", bos_token="<|im_start|>", ) chat_handler = formatter.to_chat_handler() model_path = _ensure_gguf() model = Llama( model_path=model_path, n_ctx=2048, n_gpu_layers=-1, #flash_attn=True, seed=int(time.time()), # time-based seed for varied output verbose=False, chat_handler=chat_handler, ) log.info("load_text_model: GPU model ready.") return model # ─── JSON helpers ────────────────────────────────────────────────────────────── def _extract_json(raw: str) -> dict | None: if not raw or not raw.strip(): log.warning("Model returned empty string") return None cleaned = re.sub(r".*?", "", raw, flags=re.DOTALL) cleaned = re.sub(r".*", "", cleaned, flags=re.DOTALL).strip() cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned) cleaned = re.sub(r"\s*```$", "", cleaned).strip() try: return json.loads(cleaned) except json.JSONDecodeError: pass m = re.search(r"\{.*\}", cleaned, flags=re.DOTALL) if m: try: return json.loads(m.group(0)) except json.JSONDecodeError: pass log.error("Could not parse JSON:\n%s", raw[:500]) return None def _validate_room(d: dict) -> bool: if not isinstance(d, dict): return False if not _ROOM_REQUIRED.issubset(d.keys()): log.warning("Room missing keys: %s", _ROOM_REQUIRED - d.keys()) return False if not isinstance(d.get("containers"), list) or len(d["containers"]) < 4: log.warning("Room has fewer than 4 containers") return False if not isinstance(d.get("keys"), list) or len(d["keys"]) < 2: log.warning("Room has fewer than 2 keys") return False return True # ─── Generation helpers ─────────────────────────────────────────────────────── def _chat(model, system: str, user: str, schema: dict, max_tokens: int = 1024) -> dict | None: response = model.create_chat_completion( messages=[ {"role": "system", "content": _NO_THINK_PREFIX + system}, {"role": "user", "content": user}, ], response_format={"type": "json_object", "schema": schema}, temperature=0.8, max_tokens=max_tokens, ) raw = response["choices"][0]["message"]["content"] return _extract_json(raw) def generate_game(model, max_retries: int = 3) -> dict: for attempt in range(1, max_retries + 1): data = _chat(model, GENERATE_GAME_PROMPT, '{"task":"generate_room"}', _ROOM_SCHEMA) if data and _validate_room(data): return data log.warning("generate_game attempt %d failed, retrying…", attempt) raise RuntimeError(f"generate_game failed after {max_retries} attempts.") def continue_game(model, container_name: str, key_name: str) -> str: payload = json.dumps({"container_name": container_name, "item_to_give": key_name}) data = _chat(model, CONTINUE_GAME_PROMPT, payload, _TEXT_SCHEMA) return data["text"] if data else f"You open the {container_name} and find the {key_name}." def open_door(model, room_name: str, door_name: str) -> str: payload = json.dumps({"room_name": room_name, "door_name": door_name}) data = _chat(model, OPEN_DOOR_PROMPT, payload, _TEXT_SCHEMA) return data["text"] if data else "You escape! The door swings open and freedom awaits." # ─── App ────────────────────────────────────────────────────────────────────── app = Server() _text_model = None @app.get("/") async def homepage(): html_path = Path(__file__).parent / "index.html" with open(html_path, "r", encoding="utf-8") as f: content = f.read() return HTMLResponse(content=content) @app.post("/generate_room") @spaces.GPU(duration=180) def api_generate_room(): global _text_model if _text_model is None: _text_model = load_text_model() return generate_game(_text_model) @app.post("/continue_room") @spaces.GPU(duration=120) def api_continue_room(container_name: str, key_name: str): global _text_model if _text_model is None: _text_model = load_text_model() text = continue_game(_text_model, container_name, key_name) return {"text": text} @app.post("/open_door") @spaces.GPU(duration=120) def api_open_door(room_name: str, door_name: str): global _text_model if _text_model is None: _text_model = load_text_model() text = open_door(_text_model, room_name, door_name) return {"text": text} if __name__ == "__main__": app.launch(show_error=True)