refactor: ZeroGPU-compatible — spaces first import, @spaces.GPU on generate(), deferred torch/transformers

app.py

@@ -1,25 +1,21 @@
-"""Gradio Chat UI for Aethron Portfolio Agent — CPU mode on ZeroGPU."""
+"""Gradio Chat UI for Aethron Portfolio Agent — ZeroGPU-compatible."""
+
+# ── 1. spaces MUST be imported first, before torch/transformers/faiss ──
+import spaces  # noqa: E402 — intentionally first
 
 import os
 import sys
+from threading import Thread
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
 
-import gradio as gr
-
-try:
-    import spaces
-    import spaces.zero
-    from spaces.zero import torch as spaces_torch
-    spaces_torch.patch()
-except ImportError:
-    spaces = None
+import gradio as gr  # noqa: E402 — after spaces
 
 print("=" * 55)
-print("AETHRON PORTFOLIO AGENT — CPU Mode")
+print("AETHRON PORTFOLIO AGENT — GPU Mode (ZeroGPU)")
 print("=" * 55)
 
-# Lazy-loaded inside @spaces.GPU context to avoid CUDA init in main process
+# Pipeline loaded lazily inside @spaces.GPU context — avoids CUDA init in main process
 _pipeline = None
 
 
@@ -27,52 +23,45 @@ def get_pipeline():
     global _pipeline
     if _pipeline is None:
         from rag_pipeline import AethronPipeline
-        print("Loading pipeline on CPU (first query only)...")
+        print("Loading pipeline (first query)...")
         _pipeline = AethronPipeline(
             build_index=not os.path.exists("data/index/faiss.index")
         )
     return _pipeline
 
 
-if spaces is not None:
-    @spaces.GPU(duration=120)
-    def respond(message, history):
-        if not message or not message.strip():
-            return history
-        pipeline = get_pipeline()
-        result = pipeline.query(message.strip())
-        response = result["answer"]
-        if result["sources"]:
-            source_names = []
-            for s in result["sources"][:3]:
-                name = s["section"].replace("_", " ").title()
-                if "Summary" in name:
-                    name = name.replace("Summary", "(Summary)")
-                source_names.append(name)
-            response += "\n\n**Sources:** " + " | ".join(source_names)
-        history.append({"role": "user", "content": message})
-        history.append({"role": "assistant", "content": response})
-        return history
-else:
+def _preload():
+    global _pipeline
     from rag_pipeline import AethronPipeline
+    print("Preloading pipeline in background...")
+    _pipeline = AethronPipeline(
+        build_index=not os.path.exists("data/index/faiss.index")
+    )
+    print("Pipeline preloaded and ready.")
+
+
+Thread(target=_preload, daemon=True).start()
+
 
-    def respond(message, history):
-        if not message or not message.strip():
-            return history
-        pipeline = get_pipeline()
-        result = pipeline.query(message.strip())
-        response = result["answer"]
-        if result["sources"]:
-            source_names = []
-            for s in result["sources"][:3]:
-                name = s["section"].replace("_", " ").title()
-                if "Summary" in name:
-                    name = name.replace("Summary", "(Summary)")
-                source_names.append(name)
-            response += "\n\n**Sources:** " + " | ".join(source_names)
-        history.append({"role": "user", "content": message})
-        history.append({"role": "assistant", "content": response})
+@spaces.GPU(duration=60)
+def respond(message, history):
+    """Handle chat — runs inside GPU context so torch.cuda patches apply."""
+    if not message or not message.strip():
         return history
+    pipeline = get_pipeline()
+    result = pipeline.query(message.strip())
+    response = result["answer"]
+    if result["sources"]:
+        source_names = []
+        for s in result["sources"][:3]:
+            name = s["section"].replace("_", " ").title()
+            if "Summary" in name:
+                name = name.replace("Summary", "(Summary)")
+            source_names.append(name)
+        response += "\n\n**Sources:** " + " | ".join(source_names)
+    history.append({"role": "user", "content": message})
+    history.append({"role": "assistant", "content": response})
+    return history
 
 
 CUSTOM_CSS = """
@@ -82,7 +71,7 @@ CUSTOM_CSS = """
 .header p { color: #4a4a6a; }
 """
 
-with gr.Blocks(css=CUSTOM_CSS, title="Aethron | Chat with Arash's Portfolio") as demo:
+with gr.Blocks(title="Aethron | Chat with Arash's Portfolio") as demo:
 
     gr.HTML("""
     <div class="header">
@@ -98,7 +87,7 @@ with gr.Blocks(css=CUSTOM_CSS, title="Aethron | Chat with Arash's Portfolio") as
     </div>
     """)
 
-    chatbot = gr.Chatbot(type="messages", height=500)
+    chatbot = gr.Chatbot(height=500)
 
     with gr.Row():
         msg_input = gr.Textbox(
@@ -151,4 +140,5 @@ if __name__ == "__main__":
         share=False,
         show_error=True,
         ssr_mode=False,
+        css=CUSTOM_CSS,
     )

requirements.txt

@@ -1,10 +1,10 @@
-gradio>=4.44.0
-transformers>=4.40.0
-sentence-transformers>=3.0.0
-faiss-cpu>=1.8.0
-PyMuPDF>=1.24.0
-numpy>=1.24.0
-torch>=2.3.0
-huggingface-hub>=0.24.0
-accelerate>=0.30.0
-spaces>=0.30.0
+spaces
+torch
+transformers
+sentence-transformers
+faiss-cpu
+gradio
+numpy
+huggingface-hub
+accelerate
+PyMuPDF

src/generation.py

@@ -1,11 +1,14 @@
-"""SLM generation with transformers."""
+"""SLM generation with transformers — ZeroGPU-compatible."""
+
+# ── 1. spaces MUST be imported first, before torch/transformers ──
+try:
+    import spaces  # noqa: E402 — intentionally first
+except ImportError:
+    spaces = None
 
 import os
 from typing import List, Dict
 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-
 from config import LLM_MODEL_ID, LLM_MAX_NEW_TOKENS, LLM_TEMPERATURE, LLM_DO_SAMPLE
 
 
@@ -27,10 +30,15 @@ RULES:
     TECHNICAL_PROMPT = "You are speaking with a technical partner or peer. Emphasize: architecture decisions, implementation details, repository names, and design philosophy. Be precise and technically deep."
 
     def __init__(self):
+        # Deferred imports — torch/transformers loaded inside @spaces.GPU context
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
         print(f"Loading SLM: {LLM_MODEL_ID}...")
 
         self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID)
 
+        # Model stays on CPU at init; @spaces.GPU moves it to GPU at inference time
         self.model = AutoModelForCausalLM.from_pretrained(
             LLM_MODEL_ID, dtype="auto", device_map="cpu"
         )
@@ -48,8 +56,11 @@ RULES:
 
         print("SLM loaded successfully")
 
+    @spaces.GPU(duration=60)
     def generate(self, query: str, chunks: List[Dict], persona: str = "general") -> Dict:
-        """Generate response from retrieved context."""
+        """Generate response — tokenization, forward pass, and generation inside GPU context."""
+        import torch  # noqa: F811 — re-import safe inside GPU context
+
         context_text = "\n\n".join([
             f"[Source: {c['section_type']} - {c['title']}]\n{c['text']}"
             for c in chunks