Update app.py

app.py

@@ -5,79 +5,72 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 from threading import Thread
 from duckduckgo_search import DDGS
 
-# --- CONFIG ---
-MODEL_ID = "google/gemma-3-270m-it"
+# --- STEP 1: LOAD ENV VARS ---
 HF_TOKEN = os.getenv('HF_TOKEN')
+MODEL_ID = "google/gemma-3-270m-it"
+
+print(f"--- [1/5] Initializing for {MODEL_ID} ---")
 
-# --- MODEL LOADING ---
-print("--- INITIALIZING GEMMA 3 (CPU MODE) ---")
+# --- STEP 2: LOAD TOKENIZER ---
+print("--- [2/5] Loading Tokenizer... ---")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID, 
-    device_map="cpu",          # Requires 'accelerate' to be installed
-    dtype=torch.float32,       # Updated from torch_dtype to fix deprecation
-    low_cpu_mem_usage=True, 
-    trust_remote_code=True, 
-    token=HF_TOKEN
-)
+# --- STEP 3: LOAD MODEL (MEMORY OPTIMIZED) ---
+print("--- [3/5] Materializing Model (This is where hangs usually happen)... ---")
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, 
+        device_map="cpu",
+        dtype=torch.float32, 
+        low_cpu_mem_usage=True,  # CRITICAL: Prevents RAM from spiking and hanging
+        trust_remote_code=True, 
+        token=HF_TOKEN
+    )
+    print("--- [4/5] Model Loaded Successfully! ---")
+except Exception as e:
+    print(f"FATAL ERROR DURING LOADING: {e}")
 
-# Limit CPU threads so the UI stays snappy
-torch.set_num_threads(max(1, (os.cpu_count() or 2) // 2))
+# Optimize CPU threads
+torch.set_num_threads(2)
 
 def web_search(query):
-    """Fetch live data to ground the AI's response."""
-    results = []
     try:
         with DDGS() as ddgs:
-            for r in ddgs.text(query, max_results=3):
-                results.append(f"Source: {r['href']}\nContent: {r['body']}")
-        return "\n\n".join(results) if results else "No relevant web data found."
-    except Exception as e:
-        return f"Search error: {e}"
-
-def generate_response(message, history, search_enabled, max_new_tokens, temperature):
-    context = ""
-    if search_enabled:
-        print(f"Searching for: {message}")
-        context = web_search(message)
+            return "\n\n".join([f"Source: {r['href']}\n{r['body']}" for r in ddgs.text(query, max_results=3)])
+    except:
+        return "Search failed."
 
-    # Gemma 3 Prompt Template
+def generate(message, history, search_enabled, tokens, temp):
+    context = web_search(message) if search_enabled else ""
     prompt = f"Context: {context}\n\nUser: {message}\nAssistant:"
     
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
     
-    generate_kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=int(max_new_tokens),
-        do_sample=True,
-        temperature=float(temperature),
-        top_p=0.9,
+    gen_kwargs = dict(
+        **inputs, streamer=streamer, max_new_tokens=int(tokens),
+        do_sample=True, temperature=float(temp), top_p=0.9,
     )
 
-    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
 
-    partial_text = ""
-    for new_text in streamer:
-        partial_text += new_text
-        yield partial_text
+    response = ""
+    for text in streamer:
+        response += text
+        yield response
 
-# --- GRADIO UI ---
+# --- STEP 4: UI SETUP ---
+print("--- [5/5] Launching Gradio UI... ---")
 demo = gr.ChatInterface(
-    fn=generate_response,
+    fn=generate,
     additional_inputs=[
-        gr.Checkbox(label="🌐 Enable Web Search", value=True),
-        gr.Slider(minimum=128, maximum=1024, value=512, step=128, label="Max Tokens"),
-        gr.Slider(minimum=0.1, maximum=1.2, value=0.7, step=0.1, label="Temperature"),
+        gr.Checkbox(label="Search Web", value=True),
+        gr.Slider(128, 1024, 512, label="Max Tokens"),
+        gr.Slider(0.1, 1.2, 0.7, label="Temp"),
     ],
-    title="Gemma 3 Orbit Explorer",
-    description="Optimized for CPU. This bot uses DuckDuckGo to stay up to date.",
-    theme="glass",
-    type="messages" 
+    type="messages"
 )
 
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0")
+    demo.launch(server_name="0.0.0.0")