Update app.py

app.py

@@ -8,7 +8,7 @@ from datetime import datetime
 import ast
 import operator as op
 import wikipedia
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 class Tool:
@@ -41,6 +41,10 @@ def wikipedia_search(query: str) -> str:
         wikipedia.set_lang("en")
         summary = wikipedia.summary(query, sentences=3, auto_suggest=True)
         return f"Wikipedia: {summary}"
+    except wikipedia.exceptions.DisambiguationError as e:
+        return f"Wikipedia: Multiple results found. Options: {', '.join(e.options[:5])}"
+    except wikipedia.exceptions.PageError:
+        return f"Wikipedia: No page found for '{query}'."
     except Exception as e:
         return f"Wikipedia error: {str(e)}"
 
@@ -99,36 +103,80 @@ TOOLS = [
 ]
 
 MODEL_NAME = "openai/gpt-oss-20b"
-pipe = None
+model = None
+tokenizer = None
 model_loaded = False
 
 def download_and_load_model(progress=gr.Progress()):
-    global pipe, model_loaded
+    global model, tokenizer, model_loaded
     
     try:
-        progress(0, desc="Downloading model...")
-        progress(0.5, desc="Loading model...")
+        progress(0, desc="Downloading tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
         
-        pipe = pipeline(
-            "text-generation",
-            model=MODEL_NAME,
-            torch_dtype=torch.bfloat16,
+        progress(0.4, desc="Downloading model (this may take several minutes)...")
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            trust_remote_code=True,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
             device_map="auto",
+            low_cpu_mem_usage=True,
         )
         
         progress(0.95, desc="Finalizing...")
         model_loaded = True
         progress(1.0, desc="Model loaded!")
-        return "Model loaded successfully!"
+        return f"Model '{MODEL_NAME}' loaded successfully!"
     except Exception as e:
         return f"Error: {str(e)}"
 
 def get_tool_descriptions() -> str:
     return "\n".join([f"- {tool.name}: {tool.description}" for tool in TOOLS])
 
+THINK_ONLY_PROMPT = """You are a helpful AI assistant. Solve problems step-by-step.
+
+Format:
+Thought: your reasoning
+Answer: your final answer
+
+Question: {question}
+
+Think step by step:"""
+
+ACT_ONLY_PROMPT = """You are a helpful AI assistant with tools.
+
+Available tools:
+{tools}
+
+Format:
+Action: tool_name
+Action Input: input
+
+Question: {question}
+
+Action:"""
+
+REACT_PROMPT = """You are a helpful AI assistant with tools.
+
+Available tools:
+{tools}
+
+Pattern:
+Thought: what to do next
+Action: tool_name
+Action Input: input
+Observation: [result]
+... repeat as needed
+Thought: I know the answer
+Answer: final answer
+
+Question: {question}
+
+Thought:"""
+
 def parse_action(text: str) -> tuple:
     action_match = re.search(r'Action:\s*(\w+)', text, re.IGNORECASE)
-    input_match = re.search(r'Action Input:\s*(.+?)(?=\n|$)', text, re.IGNORECASE | re.DOTALL)
+    input_match = re.search(r'Action Input:\s*(.+?)(?=\n(?:Thought:|Action:|Answer:|$))', text, re.IGNORECASE | re.DOTALL)
     return (action_match.group(1).strip(), input_match.group(1).strip()) if action_match and input_match else (None, None)
 
 def call_tool(tool_name: str, tool_input: str) -> str:
@@ -137,129 +185,161 @@ def call_tool(tool_name: str, tool_input: str) -> str:
             return tool(tool_input)
     return f"Error: Tool '{tool_name}' not found."
 
-def call_llm(messages: List[Dict], max_tokens: int = 500) -> str:
+def call_llm(prompt: str, max_tokens: int = 500) -> str:
     if not model_loaded:
-        return "Error: Model not loaded. Please click 'Load Model' button first."
+        return "Error: Model not loaded."
     
     try:
-        outputs = pipe(messages, max_new_tokens=max_tokens, do_sample=True, temperature=0.7)
-        return outputs[0]["generated_text"][-1]["content"]
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+        
+        if torch.cuda.is_available():
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                temperature=0.7,
+                do_sample=True,
+                top_p=0.9,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+        
+        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+        return response.strip()
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error during generation: {str(e)}"
 
-def chat_function(message, history, mode):
-    """Main chat function"""
-    
+def think_only_mode(question: str) -> str:
     if not model_loaded:
-        return history + [[message, "Error: Model not loaded. Please click 'Load Model' button first."]]
+        return "Error: Model not loaded."
     
-    if mode == "Think-Only":
-        system_prompt = "You are a helpful AI assistant. Think step-by-step and provide clear answers."
-        
-    elif mode == "Act-Only":
-        system_prompt = f"""You are a helpful AI assistant with tools.
-
-Available tools:
-{get_tool_descriptions()}
-
-Use tools when needed. Format:
-Action: tool_name
-Action Input: input"""
-        
-    else:
-        system_prompt = f"""You are a helpful AI assistant with tools.
-
-Available tools:
-{get_tool_descriptions()}
-
-Think step-by-step and use tools when needed. Format:
-Thought: [reasoning]
-Action: tool_name
-Action Input: input"""
+    output = "**Mode: Think-Only**\n\n"
+    response = call_llm(THINK_ONLY_PROMPT.format(question=question), max_tokens=800)
     
-    messages = [{"role": "system", "content": system_prompt}]
+    if response.startswith("Error"):
+        return output + response
     
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
+    for line in response.split('\n'):
+        if line.strip():
+            output += f"**{line.strip()}**\n\n" if line.strip().startswith(('Thought:', 'Answer:')) else f"{line}\n\n"
     
-    messages.append({"role": "user", "content": message})
+    return output + "\n---\n"
+
+def act_only_mode(question: str, max_iterations: int = 5) -> str:
+    if not model_loaded:
+        return "Error: Model not loaded."
     
-    response_parts = []
-    max_iterations = 3
+    output = "**Mode: Act-Only**\n\n"
+    conversation = ACT_ONLY_PROMPT.format(question=question, tools=get_tool_descriptions())
     
     for iteration in range(max_iterations):
-        response = call_llm(messages, max_tokens=400)
+        response = call_llm(conversation, max_tokens=300)
         
-        if mode == "ReAct":
-            thoughts = re.findall(r'Thought:\s*(.+?)(?=\nAction:|$)', response, re.IGNORECASE | re.DOTALL)
-            for thought in thoughts:
-                response_parts.append(f"💭 {thought.strip()}")
+        if response.startswith("Error"):
+            return output + response
         
-        action_name, action_input = parse_action(response)
+        output += f"**Iteration {iteration + 1}:**\n{response}\n\n"
         
-        if action_name and action_input and mode in ["Act-Only", "ReAct"]:
-            response_parts.append(f"🔧 Using: {action_name}")
+        if 'Answer:' in response:
+            match = re.search(r'Answer:\s*(.+)', response, re.IGNORECASE | re.DOTALL)
+            if match:
+                output += f"**Final Answer:** {match.group(1).strip()}\n\n"
+                break
+        
+        action_name, action_input = parse_action(response)
+        if action_name and action_input:
+            output += f"**Action:** {action_name}\n**Input:** {action_input}\n\n"
             observation = call_tool(action_name, action_input)
-            response_parts.append(f"📊 {observation}")
-            
-            messages.append({"role": "assistant", "content": response})
-            messages.append({"role": "user", "content": f"Result: {observation}\n\nProvide final answer."})
+            output += f"**Observation:** {observation}\n\n"
+            conversation += f"\n{response}\nObservation: {observation}\n\n"
         else:
-            response_parts.append(response)
+            output += f"No valid action found in iteration {iteration + 1}.\n\n"
             break
     
-    final_response = "\n\n".join(response_parts)
-    return history + [[message, final_response]]
+    return output + "\n---\n"
 
-with gr.Blocks(title="LLM Reasoning Chat") as demo:
-    gr.Markdown("""
-    # LLM Reasoning Chat
-    
-    **Model:** openai/gpt-oss-20b
-    
-    **Modes:**
-    - **Think-Only**: Pure reasoning
-    - **Act-Only**: Uses tools
-    - **ReAct**: Thinks and uses tools
-    """)
+def react_mode(question: str, max_iterations: int = 5) -> str:
+    if not model_loaded:
+        return "Error: Model not loaded."
     
-    with gr.Row():
-        load_btn = gr.Button("Load Model", variant="primary", size="lg")
-        status = gr.Textbox(label="Status", value="Click 'Load Model'", interactive=False)
+    output = "**Mode: ReAct**\n\n"
+    conversation = REACT_PROMPT.format(question=question, tools=get_tool_descriptions())
     
-    mode_selector = gr.Radio(
-        choices=["Think-Only", "Act-Only", "ReAct"],
-        value="ReAct",
-        label="Reasoning Mode"
-    )
+    for iteration in range(max_iterations):
+        response = call_llm(conversation, max_tokens=400)
+        
+        if response.startswith("Error"):
+            return output + response
+        
+        output += f"**Iteration {iteration + 1}:**\n"
+        
+        for thought in re.findall(r'Thought:\s*(.+?)(?=\n(?:Action:|Answer:|$))', response, re.IGNORECASE | re.DOTALL):
+            output += f"**Thought:** {thought.strip()}\n\n"
+        
+        if 'Answer:' in response:
+            match = re.search(r'Answer:\s*(.+)', response, re.IGNORECASE | re.DOTALL)
+            if match:
+                output += f"**Answer:** {match.group(1).strip()}\n\n"
+                break
+        
+        action_name, action_input = parse_action(response)
+        if action_name and action_input:
+            output += f"**Action:** {action_name}\n**Input:** {action_input}\n\n"
+            observation = call_tool(action_name, action_input)
+            output += f"**Observation:** {observation}\n\n"
+            conversation += f"\n{response}\nObservation: {observation}\n\nThought:"
+        else:
+            output += f"No action found in iteration {iteration + 1}.\n{response}\n\n"
+            break
     
-    chatbot = gr.Chatbot(label="Chat", height=400)
+    return output + "\n---\n"
+
+EXAMPLES = [
+    "What is 25 * 47?",
+    "What is the weather in Paris?",
+    "Who wrote 1984?",
+    "Calculate: 100 + 200",
+]
+
+def run_comparison(question: str, mode: str):
+    if mode == "Think-Only":
+        return think_only_mode(question), "", ""
+    elif mode == "Act-Only":
+        return "", act_only_mode(question), ""
+    elif mode == "ReAct":
+        return "", "", react_mode(question)
+    elif mode == "All (Compare)":
+        return think_only_mode(question), act_only_mode(question), react_mode(question)
+    return "Invalid mode.", "", ""
+
+with gr.Blocks(title="LLM Reasoning Modes") as demo:
+    gr.Markdown("# LLM Reasoning Modes Comparison\n\n**Model:** openai/gpt-oss-20b\n\n**Tools:** DuckDuckGo | Wikipedia | Weather | Calculator | Python")
     
     with gr.Row():
-        msg = gr.Textbox(label="Message", placeholder="Ask anything...", scale=4)
-        submit = gr.Button("Send", variant="primary", scale=1)
+        download_btn = gr.Button("Download & Load Model", variant="primary", size="lg")
+        model_status = gr.Textbox(label="Status", value="Click to download", interactive=False)
     
-    clear = gr.Button("Clear Chat")
-    
-    gr.Examples(
-        examples=[
-            "What is the capital of France and its weather?",
-            "Calculate: 1000 * (1.05 ** 3)",
-            "Who wrote To Kill a Mockingbird?",
-        ],
-        inputs=msg
-    )
-    
-    load_btn.click(fn=download_and_load_model, outputs=status)
+    with gr.Row():
+        with gr.Column(scale=3):
+            question_input = gr.Textbox(label="Question", lines=3)
+            mode_dropdown = gr.Dropdown(choices=["Think-Only", "Act-Only", "ReAct", "All (Compare)"], value="Think-Only", label="Mode")
+            submit_btn = gr.Button("Run", variant="primary", size="lg")
+        with gr.Column(scale=1):
+            gr.Markdown("**Examples**")
+            for idx, ex in enumerate(EXAMPLES):
+                gr.Button(f"Ex {idx+1}", size="sm").click(fn=lambda e=ex: e, outputs=question_input)
     
-    submit.click(fn=chat_function, inputs=[msg, chatbot, mode_selector], outputs=chatbot)
-    msg.submit(fn=chat_function, inputs=[msg, chatbot, mode_selector], outputs=chatbot)
+    gr.Markdown("---")
     
-    submit.click(lambda: "", outputs=msg)
-    msg.submit(lambda: "", outputs=msg)
+    with gr.Row():
+        think_output = gr.Markdown(label="Think-Only")
+        act_output = gr.Markdown(label="Act-Only")
+        react_output = gr.Markdown(label="ReAct")
     
-    clear.click(lambda: [], outputs=chatbot)
+    download_btn.click(fn=download_and_load_model, outputs=model_status)
+    submit_btn.click(fn=run_comparison, inputs=[question_input, mode_dropdown], outputs=[think_output, act_output, react_output])
 
 if __name__ == "__main__":
     demo.launch(share=True)