Update app.py

app.py

@@ -111,69 +111,28 @@ def download_and_load_model(progress=gr.Progress()):
     
     try:
         progress(0, desc="Downloading model...")
-        progress(0.5, desc="Loading model (this may take several minutes)...")
+        progress(0.5, desc="Loading model...")
         
         pipe = pipeline(
             "text-generation",
             model=MODEL_NAME,
-            torch_dtype="auto",
+            torch_dtype=torch.bfloat16,
             device_map="auto",
         )
         
         progress(0.95, desc="Finalizing...")
         model_loaded = True
         progress(1.0, desc="Model loaded!")
-        return f"Model '{MODEL_NAME}' loaded successfully!"
+        return f"Model 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(?:Thought:|Action:|Answer:|$))', text, re.IGNORECASE | re.DOTALL)
+    input_match = re.search(r'Action Input:\s*(.+?)(?=\n|$)', 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:
@@ -182,134 +141,211 @@ 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(prompt: str, max_tokens: int = 500) -> str:
+def call_llm(messages: List[Dict], max_tokens: int = 500) -> str:
     if not model_loaded:
-        return "Error: Model not loaded."
+        return "Error: Model not loaded. Please click 'Load Model' button first."
     
     try:
-        messages = [{"role": "user", "content": prompt}]
-        outputs = pipe(messages, max_new_tokens=max_tokens)
+        outputs = pipe(messages, max_new_tokens=max_tokens, do_sample=True, temperature=0.7)
         return outputs[0]["generated_text"][-1]["content"]
     except Exception as e:
         return f"Error: {str(e)}"
 
-def think_only_mode(question: str) -> str:
-    if not model_loaded:
-        return "Error: Model not loaded."
+def chat_think_only(message, history, mode):
+    """Think-Only mode: Pure reasoning without tools"""
+    
+    system_prompt = """You are a helpful AI assistant that thinks step-by-step. For each question:
+1. Think through the problem carefully
+2. Show your reasoning process
+3. Provide a clear answer
+
+Always respond conversationally and naturally."""
+    
+    messages = [{"role": "system", "content": system_prompt}]
     
-    output = "**Mode: Think-Only**\n\n"
-    response = call_llm(THINK_ONLY_PROMPT.format(question=question), max_tokens=800)
+    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"
+    response = call_llm(messages, max_tokens=800)
+    return response
 
-def act_only_mode(question: str, max_iterations: int = 5) -> str:
+def chat_act_only(message, history, mode):
+    """Act-Only mode: Use tools to answer"""
+    
     if not model_loaded:
         return "Error: Model not loaded."
     
-    output = "**Mode: Act-Only**\n\n"
-    conversation = ACT_ONLY_PROMPT.format(question=question, tools=get_tool_descriptions())
+    system_prompt = f"""You are a helpful AI assistant with access to tools. Use tools to find accurate information.
+
+Available tools:
+{get_tool_descriptions()}
+
+When you need to use a tool, respond with:
+Action: tool_name
+Action Input: your input
+
+After getting results, provide a natural answer to the user."""
     
-    for _ in range(max_iterations):
-        response = call_llm(conversation, max_tokens=300)
-        
-        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
+    messages = [{"role": "system", "content": system_prompt}]
+    
+    for user_msg, assistant_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": assistant_msg})
+    
+    messages.append({"role": "user", "content": message})
+    
+    max_iterations = 3
+    response_parts = []
+    
+    for iteration in range(max_iterations):
+        response = call_llm(messages, max_tokens=400)
         
         action_name, action_input = parse_action(response)
+        
         if action_name and action_input:
-            output += f"**Action:** {action_name}\n**Input:** {action_input}\n\n"
+            response_parts.append(f"🔧 Using tool: **{action_name}**")
             observation = call_tool(action_name, action_input)
-            output += f"**Observation:** {observation}\n\n"
-            conversation += f"\n{response}\nObservation: {observation}\n\n"
+            response_parts.append(f"📊 Result: {observation}\n")
+            
+            messages.append({"role": "assistant", "content": response})
+            messages.append({"role": "user", "content": f"Tool result: {observation}\n\nNow provide a natural answer to the user based on this information."})
         else:
-            output += f"Could not parse action.\n\n"
+            response_parts.append(response)
             break
     
-    return output + "\n---\n"
+    return "\n\n".join(response_parts)
 
-def react_mode(question: str, max_iterations: int = 5) -> str:
+def chat_react(message, history, mode):
+    """ReAct mode: Interleaved thinking and tool use"""
+    
     if not model_loaded:
         return "Error: Model not loaded."
     
-    output = "**Mode: ReAct**\n\n"
-    conversation = REACT_PROMPT.format(question=question, tools=get_tool_descriptions())
+    system_prompt = f"""You are a helpful AI assistant that thinks and uses tools. 
+
+Available tools:
+{get_tool_descriptions()}
+
+For each question:
+1. Think about what you need to do
+2. Use tools when you need information
+3. Think about the results
+4. Provide a clear answer
+
+Format for tool use:
+Thought: [your reasoning]
+Action: tool_name
+Action Input: input
+
+After tool results, think again and either use another tool or provide the final answer."""
+    
+    messages = [{"role": "system", "content": system_prompt}]
     
-    for _ in range(max_iterations):
-        response = call_llm(conversation, max_tokens=400)
-        
-        for thought in re.findall(r'Thought:\s*(.+?)(?=\n(?:Action:|Answer:|$))', response, re.IGNORECASE | re.DOTALL):
-            output += f"**Thought:** {thought.strip()}\n\n"
+    for user_msg, assistant_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": assistant_msg})
+    
+    messages.append({"role": "user", "content": message})
+    
+    max_iterations = 3
+    response_parts = []
+    
+    for iteration in range(max_iterations):
+        response = call_llm(messages, max_tokens=400)
         
-        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
+        thoughts = re.findall(r'Thought:\s*(.+?)(?=\n|Action:|$)', response, re.IGNORECASE | re.DOTALL)
+        for thought in thoughts:
+            response_parts.append(f"💭 **Thought:** {thought.strip()}")
         
         action_name, action_input = parse_action(response)
+        
         if action_name and action_input:
-            output += f"**Action:** {action_name}\n**Input:** {action_input}\n\n"
+            response_parts.append(f"🔧 **Action:** {action_name} with input: {action_input}")
             observation = call_tool(action_name, action_input)
-            output += f"**Observation:** {observation}\n\n"
-            conversation += f"\n{response}\nObservation: {observation}\n\nThought:"
+            response_parts.append(f"📊 **Result:** {observation}\n")
+            
+            messages.append({"role": "assistant", "content": response})
+            messages.append({"role": "user", "content": f"Observation: {observation}\n\nContinue thinking or provide final answer."})
         else:
-            if 'Answer:' not in response:
-                output += "No action found.\n\n"
+            if not thoughts:
+                response_parts.append(response)
             break
     
-    return output + "\n---\n"
+    return "\n\n".join(response_parts)
 
-EXAMPLES = [
-    "What is the capital of France and its weather?",
-    "Who wrote 'To Kill a Mockingbird'?",
-    "Calculate: 1000 * (1.05 ** 3)",
-    "What is Tokyo's population?",
-]
-
-def run_comparison(question: str, mode: str):
+def chat_response(message, history, mode):
+    """Route to appropriate mode"""
     if mode == "Think-Only":
-        return think_only_mode(question), "", ""
+        return chat_think_only(message, history, mode)
     elif mode == "Act-Only":
-        return "", act_only_mode(question), ""
+        return chat_act_only(message, history, mode)
     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.", "", ""
+        return chat_react(message, history, mode)
+    else:
+        return "Please select a reasoning 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.Blocks(title="LLM Reasoning Chat", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🤖 LLM Reasoning Chat
+    
+    **Model:** openai/gpt-oss-20b | **Tools:** DuckDuckGo, Wikipedia, Weather, Calculator, Python
+    
+    Select a reasoning mode and start chatting!
+    """)
     
     with gr.Row():
-        download_btn = gr.Button("Download & Load Model", variant="primary", size="lg")
-        model_status = gr.Textbox(label="Status", value="Click to download", interactive=False)
+        load_btn = gr.Button("🚀 Load Model", variant="primary", size="lg")
+        status = gr.Textbox(label="Status", value="Click 'Load Model' to start", interactive=False, scale=2)
     
     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="All (Compare)", 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)
+        mode_selector = gr.Radio(
+            choices=["Think-Only", "Act-Only", "ReAct"],
+            value="ReAct",
+            label="Reasoning Mode",
+            info="Think-Only: Pure reasoning | Act-Only: Tools only | ReAct: Think + Tools"
+        )
     
-    gr.Markdown("---")
+    chatbot = gr.Chatbot(
+        label="Chat",
+        height=500,
+        show_label=True,
+        avatar_images=(None, "https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.png")
+    )
     
     with gr.Row():
-        think_output = gr.Markdown(label="Think-Only")
-        act_output = gr.Markdown(label="Act-Only")
-        react_output = gr.Markdown(label="ReAct")
+        msg = gr.Textbox(
+            label="Message",
+            placeholder="Ask me anything...",
+            scale=4
+        )
+        submit = gr.Button("Send", variant="primary", scale=1)
+    
+    gr.Examples(
+        examples=[
+            "What is the capital of France and what's the weather there?",
+            "Who wrote 'To Kill a Mockingbird' and when was it published?",
+            "Calculate the compound interest on $1000 at 5% for 3 years",
+            "What is the population of Tokyo?",
+            "Explain quantum mechanics in simple terms",
+        ],
+        inputs=msg,
+        label="Example Questions"
+    )
+    
+    clear = gr.Button("🗑️ Clear Chat")
+    
+    load_btn.click(fn=download_and_load_model, outputs=status)
+    
+    msg.submit(fn=chat_response, inputs=[msg, chatbot, mode_selector], outputs=chatbot)
+    submit.click(fn=chat_response, inputs=[msg, chatbot, mode_selector], outputs=chatbot)
+    
+    msg.submit(lambda: "", None, msg)
+    submit.click(lambda: "", None, msg)
     
-    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])
+    clear.click(lambda: None, None, chatbot)
 
 if __name__ == "__main__":
     demo.launch(share=True)