Deploy Myco from CI

game/engine.py

@@ -9,119 +9,107 @@ import re
 import json
 import traceback
 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
 from game.catalog import load_mushrooms
 from game.state import collection_contains, mushroom_from_state, welcome_history
 
-print("[Myco] Loading model and tokenizer for google/gemma-3-1b-it on CPU...")
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-it")
 
-# Force the model to load onto the CPU thread safely
-model = AutoModelForCausalLM.from_pretrained(
-    "google/gemma-3-1b-it", 
-    dtype=torch.bfloat16
-)
-print("[Myco] Successfully loaded on CPU!")
+import os
+import torch
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 DEFAULT_MODEL_ID = "google/gemma-3-1b-it"
-# ───────────────────────────────────────────────────────────────────────────
-# CRITICAL FIX: Global Model Loading for Hugging Face ZeroGPU
-# ───────────────────────────────────────────────────────────────────────────
-MODEL_ID = os.getenv("MYCO_MODEL_ID", DEFAULT_MODEL_ID)
-HF_TOKEN = os.getenv("HF_BUILD_SMALL_HACKATHON_TOKEN")
-
-print(f"\n[Myco] Loading model and tokenizer for {MODEL_ID} at global scope...")
-try:
-    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-    _model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        token=HF_TOKEN,
-        torch_dtype=torch.bfloat16,
-        device_map=None,  # Must be None so ZeroGPU can wrap it safely
-        trust_remote_code=True
-    )
-    print(f"[Myco] Successfully loaded: {MODEL_ID}")
-except Exception as exc:
-    print(f"[Myco] Global load error: {exc}")
-    _model, _tokenizer = None, None
 
+# Global singletons
+_model = None
+_tokenizer = None
+_lock = threading.Lock() # Ensure this matches your existing lock variable name
 
 def _get_model_and_tokenizer():
-    return _model, _tokenizer
-
+    global _model, _tokenizer
+    if _model is not None and _tokenizer is not None:
+        return _model, _tokenizer
 
-def _get_pipeline():
-    model, _ = _get_model_and_tokenizer()
-    return model
+    with _lock:
+        if _model is not None and _tokenizer is not None:
+            return _model, _tokenizer
 
+        model_id = os.getenv("MYCO_MODEL_ID", "google/gemma-3-1b-it")
+        token = os.getenv("HF_BUILD_SMALL_HACKATHON_TOKEN")
 
-def _run_pipeline(messages_or_prompt): # Your exact function name here
-    global model, tokenizer
-    try:
-        # Change .to("cuda") to .to("cpu") for your tokenization
-        if isinstance(messages_or_prompt, list):
-            input_ids = tokenizer.apply_chat_template(
-                messages_or_prompt, 
-                tokenize=True, 
-                add_generation_prompt=True, 
-                return_tensors="pt"
-            ).to("cpu") # 👈 Changed to CPU
-            inputs = {"input_ids": input_ids}
-            input_len = input_ids.shape[1]
-        else:
-            inputs = tokenizer(str(messages_or_prompt), return_tensors="pt").to("cpu") # 👈 Changed to CPU
-            input_len = inputs["input_ids"].shape[1]
-
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs, 
-                max_new_tokens=256,
-                do_sample=True,
-                temperature=0.7
+        try:
+            print(f"\n[Myco] Loading model and tokenizer for {model_id}...")
+            
+            # Load tokenizer
+            _tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+            
+            # Load model strictly on CPU with bfloat16 precision
+            _model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                token=token,
+                torch_dtype=torch.bfloat16,
+                device_map=None, # CRITICAL: Must be None so ZeroGPU can hook it safely
+                trust_remote_code=True
             )
-
-        generated_tokens = outputs[0][input_len:]
-        return tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
-
-    except Exception as e:
-        print(f"Error during generation: {e}")
-        return f"ERROR: {str(e)}"
-        
+            print(f"[Myco] Successfully loaded: {model_id}")
+            return _model, _tokenizer
+        except Exception as exc:
+            print(f"[Myco] Load error: {exc}")
+            return None, None
 
 # ---------------------------------------------------------------------------
-# Template Helper
+# Pipeline loader
 # ---------------------------------------------------------------------------
-def _format_messages_for_template(system_prompt: str, user_prompt: str, history_turns: list = None) -> list:
-    raw_turns = [{"role": "user", "content": system_prompt}]
-    
-    if history_turns:
-        for entry in history_turns:
-            role = entry.get("role", "assistant")
-            if role == "assistant":
-                role = "model"
-            elif role == "system":
-                continue
-            
-            content = entry.get("content", "")
-            if isinstance(content, str) and content.strip():
-                raw_turns.append({"role": role, "content": content})
-                
-    raw_turns.append({"role": "user", "content": user_prompt})
-    
-    alternated_turns = []
-    for turn in raw_turns:
-        if alternated_turns and alternated_turns[-1]["role"] == turn["role"]:
-            alternated_turns[-1]["content"] += f"\n\n{turn['content']}"
-        else:
-            alternated_turns.append(turn)
-            
-    return alternated_turns
+def _get_pipeline():
+    """ Legacy compatibility wrapper so companion_status() and other 
+    startup hooks don't throw a NameError.
+    """
+    model, tokenizer = _get_model_and_tokenizer()
+    if model is not None and tokenizer is not None:
+        # Return the model instance so 'if pipe is not None' checks pass successfully
+        return model
+    return None
 
 # ---------------------------------------------------------------------------
-# Game Constants
+# GPU runner — ONLY this function gets the @spaces.GPU decorator.
 # ---------------------------------------------------------------------------
+@spaces.GPU
+def _run_pipeline(pipe_ignored, messages):
+    # 1. Fetch our raw model and tokenizer singletons
+    model, tokenizer = _get_model_and_tokenizer()
+    if model is None or tokenizer is None:
+        return "The forest is silent. (Model loading failed)"
+
+    # CRITICAL: Do NOT call model.to("cuda") manually here. 
+    # ZeroGPU's @spaces.GPU decorator handles the weight migration automatically.
+
+    # 2. Build the chat template structure natively
+    formatted_prompt = tokenizer.apply_chat_template(
+        messages, 
+        tokenize=False, 
+        add_generation_prompt=True
+    )
+
+    # 3. Tokenize and dynamically map inputs to the exact device the model is currently using
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+
+    # 4. Generate clean text without any configuration conflicts
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    # 5. Extract only the newly generated text tokens
+    input_length = inputs.input_ids.shape[1]
+    generated_tokens = outputs[0][input_length:]
+    
+    return tokenizer.decode(generated_tokens, skip_special_tokens=True)
+
+
 JSON_PROMPT_SUFFIX = (
     "\n\nRespond with a single JSON object only. No prose before or after it. "
     'Example: {"action":"pick","target":"Ruby Knuckle","thought":"It seems safe."}'
@@ -159,8 +147,9 @@ RARITY_CLUES = {
     "Legendary": "The whole clearing goes quiet — this mushroom hides part of the Elder Map.",
 }
 
+
 # ---------------------------------------------------------------------------
-# Utility
+# Utility: extract JSON or return text
 # ---------------------------------------------------------------------------
 def _extract_json_or_text(generated_text: str) -> str | None:
     if not generated_text:
@@ -175,11 +164,17 @@ def _extract_json_or_text(generated_text: str) -> str | None:
                 return json.dumps(parsed, separators=(",", ":"), ensure_ascii=False)
             except Exception:
                 continue
+
     return text or None
 
+
+# ---------------------------------------------------------------------------
+# Status
+# ---------------------------------------------------------------------------
 def companion_model_id():
     return os.getenv("MYCO_MODEL_ID", DEFAULT_MODEL_ID)
 
+
 def companion_status():
     pipe  = _get_pipeline()
     model = companion_model_id()
@@ -187,11 +182,13 @@ def companion_status():
         return f"🧠 Myco AI active ({model})"
     return f"⚠️ Myco AI fallback mode ({model} failed)"
 
+
 def hf_companion_status():
     return companion_status()
 
+
 # ---------------------------------------------------------------------------
-# LLM Calls
+# LLM call — single-turn
 # ---------------------------------------------------------------------------
 def _llm(prompt: str, context: dict | None = None) -> str | None:
     pipe = _get_pipeline()
@@ -214,10 +211,18 @@ def _llm(prompt: str, context: dict | None = None) -> str | None:
     score_line      = f"Player score: {ctx.get('score', 0)} spores. Health: {ctx.get('health', 3)}/3."
 
     system = f"{SYSTEM_PROMPT}\n\n{mushroom_line}\n{collection_line}\n{mystery_line}\n{score_line}"
-    messages = _format_messages_for_template(system, prompt + JSON_PROMPT_SUFFIX)
+
+    messages = [
+        {"role": "system", "content": system},
+        {"role": "user",   "content": prompt + JSON_PROMPT_SUFFIX},
+    ]
 
     try:
         outputs = _run_pipeline(pipe, messages)
+        print("========== MYCO OUTPUT ==========")
+        print(outputs)
+        print("=================================")
+
         if isinstance(outputs, list) and outputs:
             first     = outputs[0]
             generated = first.get("generated_text", "") if isinstance(first, dict) else str(first)
@@ -236,6 +241,10 @@ def _llm(prompt: str, context: dict | None = None) -> str | None:
         traceback.print_exc()
         return None
 
+
+# ---------------------------------------------------------------------------
+# LLM call — multi-turn chat
+# ---------------------------------------------------------------------------
 def _llm_with_history(history: list, user_message: str, context: dict) -> str | None:
     pipe = _get_pipeline()
     if not pipe:
@@ -259,17 +268,32 @@ def _llm_with_history(history: list, user_message: str, context: dict) -> str |
         f"Score: {ctx.get('score', 0)} spores. Health: {ctx.get('health', 3)}/3."
     )
 
-    messages = _format_messages_for_template(system, user_message, history[-6:])
+    messages = [{"role": "system", "content": system}]
+    for entry in history[-6:]:
+        role    = entry.get("role", "assistant")
+        content = entry.get("content", "")
+        if isinstance(content, str) and content.strip():
+            messages.append({"role": role, "content": content})
+    # Chat replies: no JSON forcing — Myco speaks naturally here.
+    messages.append({"role": "user", "content": user_message})
 
     try:
-        return _run_pipeline(pipe, messages)
+        reply = _run_pipeline(pipe, messages)
+    
+        print("========== MYCO OUTPUT ==========")
+        print(reply)
+        print("=================================")
+    
+        return reply
+    
     except Exception as exc:
         print(f"[Myco] Inference error: {exc}")
         traceback.print_exc()
         return None
 
+
 # ---------------------------------------------------------------------------
-# Helpers
+# Context builder
 # ---------------------------------------------------------------------------
 def _ctx(current: dict | None, collection: list) -> dict:
     count   = len(collection)
@@ -298,10 +322,15 @@ def _ctx(current: dict | None, collection: list) -> dict:
         })
     return ctx
 
+
+# ---------------------------------------------------------------------------
+# Fallbacks
+# ---------------------------------------------------------------------------
 def _fallback_discover(current: dict) -> str:
     name   = current.get("name", "something")
     rarity = current.get("rarity", "Common")
-    if current.get("name", "") in POISONOUS:
+    poison = current.get("name", "") in POISONOUS
+    if poison:
         return f"Wait — {name}! I've seen this before... something feels very wrong. Don't touch it yet."
     if rarity == "Legendary":
         return f"Oh! Oh! A {name}! The whole clearing just went silent. This is from the Elder Map!"
@@ -309,36 +338,48 @@ def _fallback_discover(current: dict) -> str:
         return f"A {name}... I can feel it humming. Something rare is here — maybe magical."
     return f"A {name}! Found near {current.get('habitat','the forest')}. Let me sense it first."
 
+
 def _fallback_pick(current: dict) -> str:
     if _is_poisonous(current):
         return "💀 That was poisonous! I tried to stop you... the forest goes dark."
     return f"Got {current.get('name','it')}! +{RARITY_SCORE.get(current.get('rarity','Common'),10)} spores!"
 
+
 def _fallback_study(current: dict) -> str:
     return f"I studied it carefully. Magic field updated. The clue: {RARITY_CLUES.get(current.get('rarity','Common'), '')}"
 
+
 def _fallback_collect(current: dict) -> str:
     return f"Added {current.get('name','it')} to the MycoDex! The pages feel warmer."
 
+
 def _fallback_whisper(current: dict) -> str:
     return "I followed the whisper... and remembered a path I've never walked. The mystery deepens."
 
+
 def _fallback_chat(current: dict | None) -> str:
     if current:
         return f"I feel something strange about {current.get('name','this')}... stay close to me."
     return "The forest is full of secrets. Move to a clearing and search — I'll watch for danger."
 
+
+# ---------------------------------------------------------------------------
+# Mushroom helpers
+# ---------------------------------------------------------------------------
 def _choose_mushroom(catalog=None):
     mushrooms = tuple(load_mushrooms() if catalog is None else catalog)
     weights   = [RARITY_WEIGHTS.get(m.rarity, 12) for m in mushrooms]
     return random.choices(mushrooms, weights=weights, k=1)[0]
 
+
 def _is_poisonous(current: dict) -> bool:
     return current.get("name", "") in POISONOUS or current.get("danger") == "Poisonous"
 
+
 def _score_value(current: dict) -> int:
     return RARITY_SCORE.get(current.get("rarity", "Common"), 10)
 
+
 def _score_collection(collection: list) -> int:
     total = 0
     for e in collection:
@@ -347,12 +388,14 @@ def _score_collection(collection: list) -> int:
         total += int(e.get("score_delta") or _score_value(e))
     return max(0, total)
 
+
 def _health(current: dict | None, collection: list) -> int:
     if current and current.get("game_over") == "Yes":
         return 0
     deaths = sum(1 for e in collection if e.get("game_over") == "Yes")
     return max(0, PLAYER_HEALTH - deaths)
 
+
 def _mystery_state(count: int) -> dict:
     chapter, next_ch = MYSTERY_CHAPTERS[0], None
     for c in MYSTERY_CHAPTERS:
@@ -370,9 +413,11 @@ def _mystery_state(count: int) -> dict:
         "mystery_next":  next_line,
     }
 
+
 def _story_event(count: int) -> dict:
     return FOREST_EVENTS[count % len(FOREST_EVENTS)]
 
+
 def _build_current(mushroom, collection: list) -> dict:
     count   = len(collection)
     current = mushroom.to_dict()
@@ -393,19 +438,24 @@ def _build_current(mushroom, collection: list) -> dict:
     current.update(_mystery_state(count))
     return current
 
+
 def _append(history: list, role: str, content: str) -> list:
     return [*history, {"role": role, "content": content}]
 
+
 def _safe_history(h) -> list:
     return list(h or welcome_history())
 
+
 def _safe_collection(c) -> list:
     return list(c or [])
 
+
 # ---------------------------------------------------------------------------
-# Public Game Actions
+# Public game actions
 # ---------------------------------------------------------------------------
 def discover_mushroom(collection=None, catalog=None):
+    """Discover a new mushroom. LLM narrates the moment."""
     coll     = _safe_collection(collection)
     mushroom = _choose_mushroom(catalog)
     current  = _build_current(mushroom, coll)
@@ -425,7 +475,10 @@ def discover_mushroom(collection=None, catalog=None):
     history = _append(welcome_history(), "assistant", reply)
     return mushroom, current, history
 
+
 def myco_reply(message=None, history=None, current=None, collection=None, position=None):
+    """Player chats with Myco. LLM responds in character with full context."""
+    print("MYCO_REPLY CALLED, message:", repr(message))
     hist  = _safe_history(history)
     coll  = _safe_collection(collection)
     clean = (message or "").strip()
@@ -439,9 +492,12 @@ def myco_reply(message=None, history=None, current=None, collection=None, positi
 
     return "", _append(hist, "user", clean) + [{"role": "assistant", "content": reply}]
 
+
 companion_reply = myco_reply
 
+
 def collect_current(current=None, collection=None, history=None):
+    """Collect mushroom into MycoDex. LLM narrates the entry."""
     coll = _safe_collection(collection)
     hist = _safe_history(history)
 
@@ -472,7 +528,9 @@ def collect_current(current=None, collection=None, history=None):
     reply = _llm(prompt, ctx) or _fallback_collect(current)
     return updated_coll, _append(hist, "assistant", reply)
 
+
 def pick_current(current=None, collection=None, history=None):
+    """Pick mushroom as game item. Poison = game over. LLM narrates dramatically."""
     coll = _safe_collection(collection)
     hist = _safe_history(history)
 
@@ -524,7 +582,9 @@ def pick_current(current=None, collection=None, history=None):
     reply = _llm(prompt, ctx) or _fallback_pick(current)
     return updated_coll, picked, _append(hist, "assistant", f"🍄 {reply}")
 
+
 def follow_whisper(current=None, collection=None, history=None):
+    """Follow the forest whisper. LLM reveals mystery fragments."""
     coll = _safe_collection(collection)
     hist = _safe_history(history)
 
@@ -560,8 +620,9 @@ def follow_whisper(current=None, collection=None, history=None):
     revealed = {**current, **mystery, "whisper_followed": "Yes"}
     return revealed, _append(hist, "assistant", f"🌌 {reply}")
 
-# ── FIX: Returns updated dictionary instead of a raw string ──
+
 def study_current(current=None, history=None):
+    """Study mushroom. LLM gives a careful field observation."""
     hist = _safe_history(history)
 
     if current is None:
@@ -572,55 +633,17 @@ def study_current(current=None, history=None):
         "Provide a concise field observation and one hint about its magic or danger."
     )
     reply = _llm(prompt, _ctx(current, [])) or _fallback_study(current)
-    
-    updated_current = dict(current)
-    updated_current["studied"] = "Yes"
-    return updated_current, _append(hist, "assistant", reply)
+    return reply, _append(hist, "assistant", reply)
 
-# ── FIX: Correct layout signature and returns single history tracking ──
-def eat_current(current=None, history=None):
-    hist = _safe_history(history)
 
+def eat_current(current=None, collection=None, history=None):
+    """Eating mushrooms raw is always blocked."""
+    hist = _safe_history(history)
     if current is None:
-        return _append(hist, "assistant", "There's nothing here to eat!")
-
-    ctx = _ctx(current, [])
-    prompt = (
-        f"The player is attempting to directly EAT the raw mushroom: {current.get('name','unknown')}. "
-        "This is highly forbidden, unsafe, and unidentified! React as Myco with absolute dynamic panic, "
-        "gently scold them in character for trying something so dangerous, and assertively tell them "
-        "they need to Study it instead of eating it!"
-    )
-    
-    fallback_reply = (
+        return collection, _append(hist, "assistant", "There's nothing here to eat!")
+    reply = (
         f"No no no! {current.get('name','That')} could be dangerous raw! "
         "I'll never let you eat an unidentified mushroom. Study it first!"
     )
-    
-    reply = _llm(prompt, ctx) or fallback_reply
-    return _append(hist, "assistant", reply)
-
-def get_myco_narrative(current=None):
-    prompt = (
-        "You are Myco, a tiny, emotional mushroom companion. "
-        "Write a whimsical, mysterious thought about the forest. "
-        "DO NOT use JSON. Do NOT use action tags. Just speak naturally."
-    )
-    
-    context = {}
-    if current:
-        context = {"name": current.get("name"), "rarity": current.get("rarity")}
-    
-    reply = _llm(prompt, context=context)
-    
-    if reply:
-        if reply.strip().startswith("{"):
-            try:
-                data = json.loads(reply)
-                return data.get("thought", "The forest feels deep today... ✨")
-            except:
-                pass
-        return reply
-
-    return "The forest feels deep today... ✨"
+    return collection, _append(hist, "assistant", reply)