Update app.py

app.py

@@ -4,71 +4,51 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
-import re
 
 from smolagents import CodeAgent, DuckDuckGoSearchTool
-from smolagents.models import OpenAIServerModel
+from smolagents.models import LocalModel
 
-SYSTEM_PROMPT = """You are a general AI assistant. Reason step by step, then finish with:
+# Define the system prompt
+SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question.
+Report your thoughts, and finish your answer concisely.
+Your answer should be a number OR as few words as possible OR a comma separated list 
+of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such 
+as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations 
+(e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, 
+apply the above rules depending of whether the element to be put in the list is a number or a string."""
 
-FINAL ANSWER: [YOUR FINAL ANSWER]
-
-Answer rules:
-- Numbers: no commas, units, or extra words. Just digits.
-- Strings: lowercase, no articles or abbreviations.
-- Lists: comma-separated, following the above.
-
-Examples:
-Q: What is 12 + 7?
-A: 12 + 7 = 19
-FINAL ANSWER: 19
-
-Q: Name three European capital cities.
-A: They are Amsterdam, Berlin, and Rome.
-FINAL ANSWER: amsterdam, berlin, rome
-
-Q: What is the square root of 81?
-A: \u221a81 = 9
-FINAL ANSWER: 9
-
-Now answer the following:
-"""
 
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-class PatchedOpenAIServerModel(OpenAIServerModel):
+# Gemini model wrapper
+class GeminiFlashModel(LocalModel):
+    def __init__(self, model_id="gemini-2.0-flash", api_url=None):
+        super().__init__(model_id=model_id, api_url=api_url)
+        self.system_prompt = SYSTEM_PROMPT
+
     def generate(self, messages, stop_sequences=None, **kwargs):
         if isinstance(messages, list):
             if not any(m["role"] == "system" for m in messages):
-                messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
+                messages = [{"role": "system", "content": self.system_prompt}] + messages
         else:
             raise TypeError("Expected 'messages' to be a list of message dicts")
 
         return super().generate(messages=messages, stop_sequences=stop_sequences, **kwargs)
 
+# Agent definition
 class MyAgent:
     def __init__(self):
-        self.model = PatchedOpenAIServerModel(model_id="gpt-4")
+        self.model = GeminiFlashModel(model_id="gemini-2.0-flash")
         self.agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=self.model)
 
     def __call__(self, question: str) -> str:
         return self.agent.run(question)
 
-def extract_final_answer(output: str) -> str:
-    if "FINAL ANSWER:" in output:
-        return output.split("FINAL ANSWER:")[-1].strip().rstrip('.')
-    return output.strip()
-
-def sanitize_answer(ans: str) -> str:
-    ans = re.sub(r'\$|%|,', '', ans)
-    ans = ans.strip().rstrip('.')
-    return ans
-
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     space_id = os.getenv("SPACE_ID")
 
     if profile:
-        username = profile.username.strip()
+        username = profile.username
         print(f"User logged in: {username}")
     else:
         print("User not logged in.")
@@ -107,9 +87,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         if not task_id or question_text is None:
             continue
         try:
-            raw_output = agent(question_text)
-            extracted = extract_final_answer(raw_output)
-            submitted_answer = sanitize_answer(extracted)
+            submitted_answer = agent(question_text)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e:
@@ -119,7 +97,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     if not answers_payload:
         return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
 
-    submission_data = {"username": username, "agent_code": agent_code, "answers": answers_payload}
+    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
     print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
     try:
         response = requests.post(submit_url, json=submission_data, timeout=60)
@@ -145,6 +123,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     except Exception as e:
         return f"An unexpected error occurred during submission: {e}", pd.DataFrame(results_log)
 
+# Gradio UI setup
 with gr.Blocks() as demo:
     gr.Markdown("# Basic Agent Evaluation Runner")
     gr.Markdown("""