Update app.py

app.py

@@ -1,7 +1,7 @@
 # app.py — ClimaMind on Hugging Face Spaces (Gradio)
-# Modes:
-#   PROVIDER=hf_model (default) -> calls HF Inference API for K2 (recommended for demo)
-#   PROVIDER=local              -> loads model with transformers (requires GPU Space)
+# Providers:
+#   PROVIDER=hf_model (default) -> calls HF Inference API (tries MODEL_ID then ALT_MODEL_ID)
+#   PROVIDER=local              -> loads model in Space (requires GPU)
 #   PROVIDER=stub               -> offline canned answers
 
 import os, time, json, random
@@ -9,9 +9,10 @@ import requests
 import gradio as gr
 
 # -------- Config --------
-PROVIDER = os.getenv("PROVIDER", "hf_model").strip()
-MODEL_ID = os.getenv("MODEL_ID", "MBZUAI-IFM/K2-Think-SFT").strip()
-HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
+PROVIDER     = os.getenv("PROVIDER", "hf_model").strip()
+MODEL_ID     = os.getenv("MODEL_ID", "LLM360/K2-Think").strip()              # default = public K2
+ALT_MODEL_ID = os.getenv("ALT_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct").strip() # fallback that works on serverless
+HF_TOKEN     = os.getenv("HF_TOKEN", "").strip()
 
 # -------- HTTP helper --------
 def _get(url, params=None, headers=None, timeout=12, retries=2, backoff=1.6):
@@ -45,21 +46,20 @@ def fetch_open_meteo(lat, lon):
     })
     return r.json()
 
-# -------- PM2.5 (Open-Meteo Air-Quality, free; replaces OpenAQ v3 which now needs a key) --------
+# -------- PM2.5 (Open-Meteo Air-Quality, free) --------
 def fetch_pm25(lat, lon):
     try:
         r = _get("https://air-quality-api.open-meteo.com/v1/air-quality", params={
             "latitude": lat, "longitude": lon, "hourly": "pm2_5", "timezone": "auto"
         }, headers={"User-Agent": "climamind-space"})
         j = r.json()
-        # take the most recent hour
         hourly = j.get("hourly", {})
         values = hourly.get("pm2_5") or []
         if values:
             return values[-1]
     except Exception:
         pass
-    return None  # graceful fallback
+    return None
 
 def fetch_factors(lat, lon):
     wx = fetch_open_meteo(lat, lon)
@@ -85,7 +85,7 @@ def heat_stress_index(temp_c, rh, wind_kmh):
     hs = (temp_c or 0) * 1.1 + (rh or 0) * 0.3 - (wind_kmh or 0) * 0.2
     return max(0, min(100, round(hs)))
 
-# -------- Prompt --------
+# -------- Prompt (escape literal braces in JSON) --------
 PROMPT = """You are ClimaMind, a climate reasoning assistant. Use ONLY the observations provided and return STRICT JSON.
 
 Location: {loc} (lat={lat}, lon={lon}), local time: {t_local}
@@ -118,17 +118,26 @@ def call_stub(_prompt:str)->str:
         "risk_badge": "Low"
     })
 
-def call_hf_model(prompt:str)->str:
+# Try HF Inference (MODEL_ID -> ALT_MODEL_ID), return (text, model_used)
+def call_hf_model(prompt:str) -> tuple[str, str]:
     from huggingface_hub import InferenceClient
-    client = InferenceClient(model=MODEL_ID, token=(HF_TOKEN or None))
-    out = client.text_generation(
-        prompt,
-        max_new_tokens=200,
-        temperature=0.1,
-        repetition_penalty=1.05,
-        do_sample=False,
-    )
-    return str(out)
+    attempts = [m for m in [MODEL_ID, ALT_MODEL_ID] if m]
+    for mid in attempts:
+        try:
+            client = InferenceClient(model=mid, token=(HF_TOKEN or None))
+            out = client.text_generation(
+                prompt,
+                max_new_tokens=200,
+                temperature=0.1,
+                repetition_penalty=1.05,
+                do_sample=False,
+            )
+            return str(out), mid
+        except Exception as e:
+            print(f"[HF_MODEL] Failed on {mid}: {repr(e)}")
+            continue
+    # If all failed, raise so we can stub
+    raise RuntimeError(f"No serverless provider available. Tried: {attempts}")
 
 _local_loaded = False
 def _ensure_local_loaded():
@@ -153,8 +162,9 @@ def _ensure_local_loaded():
     )
     _local_loaded = True
 
-def call_local(prompt:str)->str:
+def call_local(prompt:str)->tuple[str, str]:
     _ensure_local_loaded()
+    import torch  # import here to avoid dependency if not used
     if hasattr(tokenizer, "apply_chat_template"):
         messages = [{"role":"user","content":prompt}]
         inputs = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt").to(model.device)
@@ -169,7 +179,7 @@ def call_local(prompt:str)->str:
             repetition_penalty=1.05,
             eos_token_id=tokenizer.eos_token_id,
         )
-    return tokenizer.decode(out[0], skip_special_tokens=True)
+    return tokenizer.decode(out[0], skip_special_tokens=True), MODEL_ID
 
 def reason_answer(loc, coords, factors, query):
     d_idx = drying_index(factors.get("temp_c"), factors.get("rh"), factors.get("wind_kmh"))
@@ -183,27 +193,35 @@ def reason_answer(loc, coords, factors, query):
     )
 
     if PROVIDER == "hf_model":
-        raw = call_hf_model(prompt)
+        try:
+            raw, model_used = call_hf_model(prompt)
+        except Exception as e:
+            print("[HF_MODEL] Falling back to stub:", repr(e))
+            raw, model_used = call_stub(prompt), "stub"
     elif PROVIDER == "local":
-        raw = call_local(prompt)
+        raw, model_used = call_local(prompt)
     else:
-        raw = call_stub(prompt)
+        raw, model_used = call_stub(prompt), "stub"
 
+    # Extract JSON
     start, end = raw.find("{"), raw.rfind("}")
     if start == -1 or end == -1:
-        return {
+        parsed = {
             "answer": "The reasoning service returned non-JSON text. Please try again.",
-            "why_trace": ["Response formatting issue", "Low temperature helps", "Retry the query"],
-            "risk_badge": "Low"
-        }
-    try:
-        return json.loads(raw[start:end+1])
-    except Exception:
-        return {
-            "answer": "Failed to parse JSON from model output.",
-            "why_trace": ["JSON parsing error", "Reduce tokens/temperature", "Retry once"],
+            "why_trace": ["Response formatting issue", "Keep temperature low", "Retry once"],
             "risk_badge": "Low"
         }
+    else:
+        try:
+            parsed = json.loads(raw[start:end+1])
+        except Exception:
+            parsed = {
+                "answer": "Failed to parse JSON from model output.",
+                "why_trace": ["JSON parsing error", "Reduce tokens/temperature", "Retry once"],
+                "risk_badge": "Low"
+            }
+    parsed["_model_used"] = model_used
+    return parsed
 
 # -------- Gradio UI --------
 def app(city, question):
@@ -216,11 +234,13 @@ def app(city, question):
     fx = ", ".join([f"{k}={v}" for k, v in data["factors"].items()])
     why_list = ans.get("why_trace") or []
     why = "\n• " + "\n• ".join(why_list) if why_list else "\n• (no trace returned)"
+    model_used = ans.pop("_model_used", "unknown")
     md = (
         f"**Answer:** {ans.get('answer','(no answer)')}\n\n"
         f"**Why-trace:**{why}\n\n"
         f"**Risk:** {ans.get('risk_badge','N/A')}\n\n"
-        f"**Factors:** {fx}"
+        f"**Factors:** {fx}\n\n"
+        f"<sub>Provider: {PROVIDER} • Model: `{model_used}`</sub>"
     )
     return md
 
@@ -240,8 +260,8 @@ demo = gr.Interface(
     ],
     outputs=gr.Markdown(label="ClimaMind"),
     title="ClimaMind — K2-Think + Live Climate Data",
-    description="Provider = hf_model (Inference API) | local (GPU Space) | stub (offline). Configure env in Space settings.",
-    allow_flagging="never",
+    description="Serverless tries K2, falls back to Qwen if needed; or run locally on GPU Space. Stub as last resort.",
+    flagging_mode="never",
     concurrency_limit=2,
 )