Update app.py

app.py

@@ -2,34 +2,60 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
-from smolagents import CodeAgent, DuckDuckGoSearchTool, TransformersModel
+from smolagents import CodeAgent, DuckDuckGoSearchTool
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
+# Create our own model wrapper that handles the chat template properly
+class CustomTransformersModel:
+    def __init__(self, model_id="gpt2"):
+        self.model_id = model_id
+        # Create the tokenizer and explicitly set the chat template
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        
+        # Set the chat template directly on the tokenizer
+        simple_template = "{% for message in messages %}\n{% if message['role'] == 'user' %}\nUser: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}\nAssistant: {{ message['content'] }}\n{% elif message['role'] == 'system' %}\nSystem: {{ message['content'] }}\n{% endif %}\n{% endfor %}\n{% if add_generation_prompt %}\nAssistant: {% endif %}"
+        self.tokenizer.chat_template = simple_template
+        
+        # Load the model
+        self.model = AutoModelForCausalLM.from_pretrained(model_id)
+        
+    def __call__(self, prompt, **kwargs):
+        # Format the prompt using our chat template
+        messages = [{"role": "user", "content": prompt}]
+        formatted_prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
+        
+        # Tokenize the prompt
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt")
+        
+        # Generate the response
+        outputs = self.model.generate(
+            inputs.input_ids,
+            max_new_tokens=256,
+            do_sample=True,
+            temperature=0.7,
+            **kwargs
+        )
+        
+        # Decode the response
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract just the assistant's response
+        try:
+            assistant_response = response.split("Assistant: ")[-1]
+        except:
+            assistant_response = response
+            
+        return assistant_response
 
 # --- Define Agent ---
 class SmolAgentWrapper:
     def __init__(self):
-        # Use a model that's compatible with AutoModelForCausalLM
-        # GPT-2 should work, but we need to properly handle the chat template issue
-        self.model = TransformersModel(
-            model_id="gpt2",
-            generation_kwargs={
-                "do_sample": True,
-                "max_new_tokens": 256,
-                "temperature": 0.7,
-                "chat_template": "{% for message in messages %}\n{% if message['role'] == 'user' %}\nUser: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}\nAssistant: {{ message['content'] }}\n{% elif message['role'] == 'system' %}\nSystem: {{ message['content'] }}\n{% endif %}\n{% endfor %}\n{% if add_generation_prompt %}\nAssistant: {% endif %}"
-            }
-        )
-        
-        # Alternative options if the above doesn't work:
-        # Option 1: Using a different GPT model that might handle chat better
-        # self.model = TransformersModel(model_id="facebook/opt-350m")
-        
-        # Option 2: Using a model with better instruction following
-        # self.model = TransformersModel(model_id="databricks/dolly-v2-3b")
-        
+        # Use our custom model wrapper
+        self.model = CustomTransformersModel(model_id="gpt2")
         self.tools = [DuckDuckGoSearchTool()]
         self.agent = CodeAgent(model=self.model, tools=self.tools)
 
@@ -116,14 +142,14 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
 # --- Gradio Interface ---
 with gr.Blocks() as demo:
-    gr.Markdown("# SmolAgent Evaluation Runner")
+    gr.Markdown("# SmolAgent Evaluation Runner (Custom GPT-2 Implementation)")
     gr.Markdown(
         """
         **Instructions:**
         1. Log in to Hugging Face with the button below.
         2. Click the button to run all GAIA questions through the SmolAgent.
         3. Results will be submitted automatically and your score will be shown.
-        **Note:** Model runs on Hugging Face Inference API.
+        **Note:** Using a custom implementation to handle chat templates properly.
         """
     )