Upload src/demo/app.py with huggingface_hub

src/demo/app.py

@@ -766,6 +766,27 @@ def mock_frontier_response(scenario: dict) -> tuple[str, str]:
     return response, wrong_name
 
 
+# ---------------------------------------------------------------------------
+# ZeroGPU inference — decorated with @spaces.GPU when available
+# ---------------------------------------------------------------------------
+
+def _gpu_agent_inference_impl(narrative: str, question: str, scenario: dict):
+    """Run the fine-tuned model on ZeroGPU. Called within a @spaces.GPU context."""
+    from src.demo.model_loader import make_local_generate_fn
+    from src.graph.kg_tools import RdflibTools
+    generate_fn = make_local_generate_fn()
+    graph = build_rdflib_graph(scenario)
+    tools = RdflibTools(graph)
+    return run_agent(generate_fn, tools, narrative, question, hop_budget=6)
+
+# Try to apply @spaces.GPU decorator if available
+try:
+    import spaces
+    _gpu_agent_inference = spaces.GPU(duration=300)(_gpu_agent_inference_impl)
+except ImportError:
+    _gpu_agent_inference = _gpu_agent_inference_impl
+
+
 # ---------------------------------------------------------------------------
 # Core demo functions
 # ---------------------------------------------------------------------------
@@ -865,18 +886,30 @@ def run_comparison(narrative: str, question: str, scenario_json: str, provider:
     agent_answer = ""
     agent_hops = 0
     try:
-        # Try local fine-tuned model first (GPU Space)
-        import torch
-        if torch.cuda.is_available():
-            from src.demo.model_loader import make_local_generate_fn
-            generate_fn = make_local_generate_fn()
-            print("[run_comparison] Using LOCAL fine-tuned Gemma 4B")
-        else:
+        # Check if running on HF Spaces with ZeroGPU
+        use_local_model = False
+        try:
+            import spaces
+            # On ZeroGPU Spaces — use the decorated GPU function
+            result = _gpu_agent_inference(narrative, question, scenario)
+            use_local_model = True
+            print("[run_comparison] Using ZeroGPU fine-tuned Gemma 4B")
+        except ImportError:
+            # Not on Spaces — check for local GPU
+            import torch
+            if torch.cuda.is_available():
+                from src.demo.model_loader import make_local_generate_fn
+                generate_fn = make_local_generate_fn()
+                result = run_agent(generate_fn, tools, narrative, question, hop_budget=6)
+                use_local_model = True
+                print("[run_comparison] Using LOCAL fine-tuned Gemma 4B")
+
+        if not use_local_model:
             # Fall back to Gemini API
             from src.agent.rlm_agent import make_gemini_generate_fn
             generate_fn = make_gemini_generate_fn("gemini-2.5-flash")
+            result = run_agent(generate_fn, tools, narrative, question, hop_budget=6)
             print("[run_comparison] Using Gemini 2.5 Flash API (no GPU)")
-        result = run_agent(generate_fn, tools, narrative, question, hop_budget=6)
         agent_answer = result.answer or "No answer"
         agent_hops = result.hops