Update app.py

app.py

@@ -1,55 +1,105 @@
 import os
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from duckduckgo_search import DDGS
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from ddgs import DDGS # Updated package name
+from threading import Thread
 
 # --- CONFIG ---
 MODEL_ID = "google/gemma-3-270m-it"
 HF_TOKEN = os.getenv('HF_TOKEN')
 
-print("--- LOADING GEMMA 3 WITH SEARCH ABILITIES ---")
+# --- MODEL LOADING ---
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID, device_map="cpu", dtype="auto", low_cpu_mem_usage=True, trust_remote_code=True, token=HF_TOKEN
+    MODEL_ID, 
+    device_map="cpu", 
+    torch_dtype=torch.float32, # CPU is more stable with float32
+    low_cpu_mem_usage=True, 
+    trust_remote_code=True, 
+    token=HF_TOKEN
 )
 
-def get_web_context(query):
-    """Fetch the top 3 search results from DuckDuckGo."""
+def web_search(query):
     results = []
     try:
         with DDGS() as ddgs:
             for r in ddgs.text(query, max_results=3):
                 results.append(f"Source: {r['href']}\nContent: {r['body']}")
     except Exception as e:
-        print(f"Search error: {e}")
+        return f"Search failed: {e}"
     return "\n\n".join(results)
 
-def chat_with_search(message, history):
-    # 1. Get real-time info
-    print(f"Searching the web for: {message}")
-    web_data = get_web_context(message)
+def stream_chat(message, history, search_enabled, max_tokens, temperature):
+    # 1. Handle Web Search
+    context = ""
+    if search_enabled:
+        context = f"\n\nWEB SEARCH RESULTS:\n{web_search(message)}"
     
-    # 2. Construct a 'RAG' prompt (Retrieval-Augmented Generation)
-    prompt = f"""
-    You are an AI assistant with web access. 
-    Use the following search results to answer the user's question accurately.
+    # 2. Prepare Prompt
+    full_prompt = f"Context: {context}\n\nUser: {message}\nAssistant:"
+    inputs = tokenizer(full_prompt, return_tensors="pt").to("cpu")
     
-    SEARCH RESULTS:
-    {web_data}
+    # 3. Setup Streamer (This fixes the "Freezing" issue)
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
     
-    USER QUESTION: {message}
-    ANSWER:"""
+    generate_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=0.9,
+    )
 
-    inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_new_tokens=2048)
+    # Run generation in a separate thread so UI stays responsive
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    # Yield tokens one by one
+    partial_message = ""
+    for new_token in streamer:
+        partial_message += new_token
+        yield partial_message
+
+# --- GRADIO UI ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚀 Gemma 3 Ultra Bot (CPU Optimized)")
     
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Extract only the final answer
-    return response.split("ANSWER:")[-1].strip()
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(height=500)
+            msg = gr.Textbox(placeholder="Ask me anything...", label="Input")
+            with gr.Row():
+                submit = gr.Button("Send", variant="primary")
+                clear = gr.Button("Clear Chat")
+
+        with gr.Column(scale=1):
+            gr.Markdown("### ⚙️ Settings")
+            search_toggle = gr.Checkbox(label="Enable Web Search", value=False)
+            token_slider = gr.Slider(minimum=64, maximum=1024, value=256, step=64, label="Max New Tokens")
+            temp_slider = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature")
+            gr.Markdown("---")
+            gr.Info("Note: Generation on CPU may take 10-30 seconds. Streaming is enabled to show progress.")
+
+    # Link components
+    def user(user_message, history):
+        return "", history + [[user_message, None]]
+
+    def bot(history, search_on, tokens, temp):
+        user_message = history[-1][0]
+        history[-1][1] = ""
+        for character in stream_chat(user_message, history, search_on, tokens, temp):
+            history[-1][1] = character
+            yield history
+
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, search_toggle, token_slider, temp_slider], chatbot
+    )
+    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, search_toggle, token_slider, temp_slider], chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
 
-# Launching
-demo = gr.ChatInterface(fn=chat_with_search)
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", share=True)
+    demo.queue().launch(server_name="0.0.0.0")