Update app.py

app.py

@@ -11,7 +11,7 @@ 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"):
+    def __init__(self, model_id="EleutherAI/gpt-neo-125m"):
         self.model_id = model_id
         # Create the tokenizer and explicitly set the chat template
         self.tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -50,12 +50,43 @@ class CustomTransformersModel:
             assistant_response = response
             
         return assistant_response
+        
+    # Add generate method to match the interface expected by CodeAgent
+    def generate(self, prompt, **kwargs):
+        return self(prompt, **kwargs)
+    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 our custom model wrapper
-        self.model = CustomTransformersModel(model_id="gpt2")
+        # Use our custom model wrapper with GPT-Neo
+        self.model = CustomTransformersModel(model_id="EleutherAI/gpt-neo-125m")
         self.tools = [DuckDuckGoSearchTool()]
         self.agent = CodeAgent(model=self.model, tools=self.tools)
 
@@ -142,14 +173,14 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
 # --- Gradio Interface ---
 with gr.Blocks() as demo:
-    gr.Markdown("# SmolAgent Evaluation Runner (Custom GPT-2 Implementation)")
+    gr.Markdown("# SmolAgent Evaluation Runner (GPT-Neo 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:** Using a custom implementation to handle chat templates properly.
+        **Note:** Using GPT-Neo 125M with custom chat template implementation.
         """
     )