Update app.py

app.py

@@ -5,51 +5,56 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 from threading import Thread
 from duckduckgo_search import DDGS
 
-# --- STEP 1: LOAD ENV VARS ---
+# --- CONFIG ---
 HF_TOKEN = os.getenv('HF_TOKEN')
 MODEL_ID = "google/gemma-3-270m-it"
 
-print(f"--- [1/5] Initializing for {MODEL_ID} ---")
-
-# --- STEP 2: LOAD TOKENIZER ---
-print("--- [2/5] Loading Tokenizer... ---")
+print("--- [1/5] Initializing ---")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, 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}")
+print("--- [2/5] Loading Model ---")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, 
+    device_map="cpu",
+    dtype=torch.float32, 
+    low_cpu_mem_usage=True, 
+    trust_remote_code=True, 
+    token=HF_TOKEN
+)
+print("--- [3/5] Model Loaded! ---")
 
-# Optimize CPU threads
+# Optimize for CPU
 torch.set_num_threads(2)
 
 def web_search(query):
     try:
         with DDGS() as ddgs:
-            return "\n\n".join([f"Source: {r['href']}\n{r['body']}" for r in ddgs.text(query, max_results=3)])
+            results = [f"Source: {r['href']}\n{r['body']}" for r in ddgs.text(query, max_results=3)]
+            return "\n\n".join(results)
     except:
-        return "Search failed."
+        return "Search currently unavailable."
 
 def generate(message, history, search_enabled, tokens, temp):
-    context = web_search(message) if search_enabled else ""
-    prompt = f"Context: {context}\n\nUser: {message}\nAssistant:"
+    # In older Gradio, history is a list of lists: [[user, bot], [user, bot]]
+    # We just need the current message and the search toggle
+    
+    context = ""
+    if search_enabled:
+        print(f"Searching web for: {message}")
+        context = web_search(message)
+    
+    prompt = f"System: Use context to help.\nContext: {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)
     
     gen_kwargs = dict(
-        **inputs, streamer=streamer, max_new_tokens=int(tokens),
-        do_sample=True, temperature=float(temp), top_p=0.9,
+        **inputs, 
+        streamer=streamer, 
+        max_new_tokens=int(tokens),
+        do_sample=True, 
+        temperature=float(temp), 
+        top_p=0.9,
     )
 
     thread = Thread(target=model.generate, kwargs=gen_kwargs)
@@ -60,17 +65,22 @@ def generate(message, history, search_enabled, tokens, temp):
         response += text
         yield response
 
-# --- STEP 4: UI SETUP ---
-print("--- [5/5] Launching Gradio UI... ---")
+print("--- [4/5] Building Interface ---")
+
+# Removed 'type' argument to ensure compatibility with Gradio 4
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
-        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"),
+        gr.Checkbox(label="Enable Web Search", value=True),
+        gr.Slider(128, 1024, 512, step=64, label="Max New Tokens"),
+        gr.Slider(0.1, 1.2, 0.7, step=0.1, label="Temperature"),
     ],
-    type="messages"
+    title="Gemma 3 Web Search Bot",
+    theme="soft"
 )
 
+print("--- [5/5] Launching! ---")
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0")
+    # If share=True fails, OrbitMC might not allow tunnels. 
+    # Try with it first, then remove if it crashes.
+    demo.launch(server_name="0.0.0.0", share=True)