Update app.py

app.py

@@ -8,9 +8,8 @@ from datetime import datetime
 import ast
 import operator as op
 import wikipedia
-from huggingface_hub import snapshot_download
 import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 class Tool:
     def __init__(self, name: str, description: str, func):
@@ -175,31 +174,20 @@ MODEL_NAME = "openai/gpt-oss-20b"
 model = None
 tokenizer = None
 model_loaded = False
-model_path = None
 
 def download_and_load_model(progress=gr.Progress()):
     """Download and load the model."""
-    global model, tokenizer, model_loaded, model_path
+    global model, tokenizer, model_loaded
     
     try:
         progress(0, desc="Starting model download...")
         
-        progress(0.1, desc="Downloading model files (this will take several minutes)...")
-        model_path = snapshot_download(
-            repo_id=MODEL_NAME,
-            cache_dir="./model_cache",
-            resume_download=True
-        )
-        
-        progress(0.6, desc="Loading tokenizer...")
-        tokenizer = GPT2Tokenizer.from_pretrained(model_path)
-        
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
+        progress(0.3, desc="Downloading tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
         
-        progress(0.7, desc="Loading model into memory...")
-        model = GPT2LMHeadModel.from_pretrained(
-            model_path,
+        progress(0.5, desc="Downloading model (this may take several minutes)...")
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
             device_map="auto" if torch.cuda.is_available() else None,
             low_cpu_mem_usage=True
@@ -209,7 +197,7 @@ def download_and_load_model(progress=gr.Progress()):
         model_loaded = True
         
         progress(1.0, desc="Model loaded successfully!")
-        return f"Model 'openai/gpt-oss-20b' loaded successfully! Model path: {model_path}"
+        return f"Model '{MODEL_NAME}' loaded successfully!"
     
     except Exception as e:
         return f"Error loading model: {str(e)}"
@@ -290,7 +278,7 @@ def call_tool(tool_name: str, tool_input: str) -> str:
             return tool(tool_input)
     return f"Error: Tool '{tool_name}' not found. Available tools: {', '.join([t.name for t in TOOLS])}"
 
-def call_llm(messages: List[Dict], temperature: float = 0.7, max_tokens: int = 500) -> str:
+def call_llm(prompt: str, temperature: float = 0.7, max_tokens: int = 500) -> str:
     """Call the local LLM."""
     global model, tokenizer, model_loaded
     
@@ -298,11 +286,15 @@ def call_llm(messages: List[Dict], temperature: float = 0.7, max_tokens: int = 5
         return "Error: Model not loaded. Please click 'Download & Load Model' first."
     
     try:
-        prompt = messages[0]["content"]
+        messages = [{"role": "user", "content": prompt}]
         
-        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()}
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(model.device)
         
         with torch.no_grad():
             outputs = model.generate(
@@ -310,11 +302,10 @@ def call_llm(messages: List[Dict], temperature: float = 0.7, max_tokens: int = 5
                 max_new_tokens=max_tokens,
                 temperature=temperature,
                 do_sample=True,
-                top_p=0.9,
-                pad_token_id=tokenizer.eos_token_id
+                top_p=0.9
             )
         
-        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+        response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
         return response.strip()
     
     except Exception as e:
@@ -326,12 +317,11 @@ def think_only_mode(question: str) -> str:
         return "Error: Model not loaded. Please click 'Download & Load Model' first."
     
     prompt = THINK_ONLY_PROMPT.format(question=question)
-    messages = [{"role": "user", "content": prompt}]
     
     output = "**Mode: Think-Only (Chain-of-Thought)**\n\n"
     output += "Generating thoughts...\n\n"
     
-    response = call_llm(messages, temperature=0.7, max_tokens=800)
+    response = call_llm(prompt, temperature=0.7, max_tokens=800)
     
     lines = response.split('\n')
     for line in lines:
@@ -356,13 +346,13 @@ def act_only_mode(question: str, max_iterations: int = 5) -> str:
     
     output = "**Mode: Act-Only (Tool Use Only)**\n\n"
     
-    messages = [{"role": "user", "content": prompt}]
     iteration = 0
+    conversation_history = prompt
     
     while iteration < max_iterations:
         iteration += 1
         
-        response = call_llm(messages, temperature=0.5, max_tokens=300)
+        response = call_llm(conversation_history, temperature=0.5, max_tokens=300)
         
         if 'Answer:' in response:
             answer_match = re.search(r'Answer:\s*(.+)', response, re.IGNORECASE | re.DOTALL)
@@ -379,8 +369,7 @@ def act_only_mode(question: str, max_iterations: int = 5) -> str:
             observation = call_tool(action_name, action_input)
             output += f"**Observation:** {observation}\n\n"
             
-            messages.append({"role": "assistant", "content": response})
-            messages.append({"role": "user", "content": f"Observation: {observation}\n\nContinue with another action or provide the final answer."})
+            conversation_history += f"\n{response}\nObservation: {observation}\n\nContinue with another action or provide the final answer.\n"
         else:
             output += f"Could not parse action from response. Response: {response}\n\n"
             break
@@ -401,13 +390,13 @@ def react_mode(question: str, max_iterations: int = 5) -> str:
     
     output = "**Mode: ReAct (Thought + Action + Observation)**\n\n"
     
-    messages = [{"role": "user", "content": prompt}]
     iteration = 0
+    conversation_history = prompt
     
     while iteration < max_iterations:
         iteration += 1
         
-        response = call_llm(messages, temperature=0.7, max_tokens=400)
+        response = call_llm(conversation_history, temperature=0.7, max_tokens=400)
         
         thought_matches = re.findall(r'Thought:\s*(.+?)(?=\n(?:Action:|Answer:|$))', response, re.IGNORECASE | re.DOTALL)
         for thought in thought_matches:
@@ -428,8 +417,7 @@ def react_mode(question: str, max_iterations: int = 5) -> str:
             observation = call_tool(action_name, action_input)
             output += f"**Observation:** {observation}\n\n"
             
-            messages.append({"role": "assistant", "content": response})
-            messages.append({"role": "user", "content": f"Observation: {observation}\n\nThought:"})
+            conversation_history += f"\n{response}\nObservation: {observation}\n\nThought:"
         else:
             if 'Answer:' not in response:
                 output += f"No action found. Response: {response}\n\n"
@@ -464,17 +452,6 @@ def run_comparison(question: str, mode: str):
         return "Invalid mode selected.", "", ""
 
 with gr.Blocks(title="LLM Reasoning Modes Comparison") as demo:
-    gr.Markdown("""
-    # LLM Reasoning Modes Comparison
-    
-    Compare three reasoning approaches using **openai/gpt-oss-20b**:
-    
-    - **Think-Only**: Chain-of-Thought reasoning only (no tools)
-    - **Act-Only**: Tool use only (no explicit reasoning)
-    - **ReAct**: Interleaved Thought, Action, Observation
-    
-    **Available Tools:** DuckDuckGo Search | Wikipedia | Weather API | Calculator | Python REPL
-    """)
     
     with gr.Row():
         download_btn = gr.Button("Download & Load Model", variant="primary", size="lg")