Update app.py

app.py

@@ -8,8 +8,9 @@ from datetime import datetime
 import ast
 import operator as op
 import wikipedia
+from huggingface_hub import snapshot_download
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
 class Tool:
     def __init__(self, name: str, description: str, func):
@@ -170,37 +171,45 @@ TOOLS = [
     ),
 ]
 
-MODEL_NAME = "meta-llama/Llama-3.2-3B-Instruct"
+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
+    global model, tokenizer, model_loaded, model_path
     
     try:
         progress(0, desc="Starting model download...")
         
-        progress(0.2, desc="Downloading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        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.4, desc="Downloading model (this may take several minutes)...")
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
+        progress(0.7, desc="Loading model into memory...")
+        model = GPT2LMHeadModel.from_pretrained(
+            model_path,
             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
         )
         
-        progress(0.9, desc="Finalizing...")
+        progress(0.95, desc="Finalizing...")
         model_loaded = True
         
         progress(1.0, desc="Model loaded successfully!")
-        return f"Model '{MODEL_NAME}' loaded successfully! You can now run queries."
+        return f"Model 'openai/gpt-oss-20b' loaded successfully! Model path: {model_path}"
     
     except Exception as e:
         return f"Error loading model: {str(e)}"
@@ -458,7 +467,7 @@ with gr.Blocks(title="LLM Reasoning Modes Comparison") as demo:
     gr.Markdown("""
     # LLM Reasoning Modes Comparison
     
-    Compare three reasoning approaches using **Llama-3.2-3B-Instruct**:
+    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)
@@ -469,7 +478,7 @@ with gr.Blocks(title="LLM Reasoning Modes Comparison") as demo:
     
     with gr.Row():
         download_btn = gr.Button("Download & Load Model", variant="primary", size="lg")
-        model_status = gr.Textbox(label="Model Status", value="Model not loaded. Click the button to download.", interactive=False)
+        model_status = gr.Textbox(label="Model Status", value="Model not loaded. Click to download openai/gpt-oss-20b", interactive=False)
     
     with gr.Row():
         with gr.Column(scale=3):