Update app.py

app.py

@@ -1,6 +1,4 @@
 import os
-import math
-import datetime
 import requests
 import pandas as pd
 import gradio as gr
@@ -8,61 +6,73 @@ import torch
 
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from smolagents import CodeAgent, DuckDuckGoSearchTool
-from smolagents.tools import Tool
 
-DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-MODEL_ID = "google/flan-t5-xl"
+DEFAULT_API_URL = os.getenv("DEFAULT_API_URL", "https://agents-course-unit4-scoring.hf.space")
 
-# DuckDuckGo-based Wikipedia-style tool
-class WikipediaSearchTool(DuckDuckGoSearchTool, Tool):
+# IMPORTANT: xl will likely OOM on CPU/free Spaces.
+# Override via Space variables if you have a beefier GPU Space.
+DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "google/flan-t5-base")
+
+
+# DuckDuckGo-based Wikipedia-style tool (just rename + description)
+class WikipediaSearchTool(DuckDuckGoSearchTool):
     name = "wikipedia_search"
     description = "Search Wikipedia-style information using DuckDuckGo backend."
 
-# Local Flan-T5 model wrapper
+
 class LocalFlanT5Model:
-    def __init__(self, model_id=MODEL_ID):
+    """
+    Minimal wrapper compatible with smolagents' model interface:
+    model(prompt) -> str
+    """
+    def __init__(self, model_id: str = DEFAULT_MODEL_ID):
+        self.model_id = model_id
         self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model.to(self.device)
 
-    def _sanitize_prompt(self, raw_prompt):
+        # Be conservative: CPU uses fp32, CUDA uses fp16 if possible.
+        if self.device == "cuda":
+            self.model = AutoModelForSeq2SeqLM.from_pretrained(
+                model_id,
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=True,
+            ).to(self.device)
+        else:
+            self.model = AutoModelForSeq2SeqLM.from_pretrained(
+                model_id,
+                torch_dtype=torch.float32,
+                low_cpu_mem_usage=True,
+            ).to(self.device)
+
+        self.model.eval()
+
+    def _sanitize_prompt(self, raw_prompt) -> str:
+        # smolagents may pass: str OR [{"role": "...", "content": "..."}]
         if isinstance(raw_prompt, str):
             return raw_prompt.strip()
 
-        if isinstance(raw_prompt, dict):
-            for key in ["prompt", "text", "question", "input"]:
-                value = raw_prompt.get(key)
-                if isinstance(value, str):
-                    return value.strip()
-                elif isinstance(value, list):
-                    return " ".join([str(v) for v in value if isinstance(v, str)])
-
         if isinstance(raw_prompt, list):
             parts = []
             for item in raw_prompt:
                 if isinstance(item, dict):
-                    content = item.get("content")
-                    if isinstance(content, str):
-                        parts.append(content)
+                    c = item.get("content")
+                    if isinstance(c, str):
+                        parts.append(c)
                 elif isinstance(item, str):
                     parts.append(item)
-                else:
-                    try:
-                        content = getattr(item, "content", None)
-                        if isinstance(content, str):
-                            parts.append(content)
-                    except Exception:
-                        continue
-            return " ".join(parts).strip()
-
-        # Fallback
-        try:
-            return str(raw_prompt).strip()
-        except Exception:
-            raise ValueError(f"Prompt must be a string, got {type(raw_prompt)}: {raw_prompt}")
+            return "\n".join(parts).strip()
+
+        if isinstance(raw_prompt, dict):
+            for key in ("prompt", "text", "question", "input", "content"):
+                v = raw_prompt.get(key)
+                if isinstance(v, str) and v.strip():
+                    return v.strip()
+
+        return str(raw_prompt).strip()
 
-    def generate(self, prompt, **kwargs) -> str:
+    @torch.inference_mode()
+    def generate(self, prompt, max_new_tokens: int = 256) -> str:
         prompt_str = self._sanitize_prompt(prompt)
         if not prompt_str:
             raise ValueError("Prompt must be a non-empty string.")
@@ -71,54 +81,75 @@ class LocalFlanT5Model:
             prompt_str,
             return_tensors="pt",
             truncation=True,
-            max_length=self.tokenizer.model_max_length
+            max_length=min(self.tokenizer.model_max_length, 512),
         ).to(self.device)
 
-        outputs = self.model.generate(**inputs, max_new_tokens=500)
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        outputs = self.model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+        )
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
     def __call__(self, prompt):
         return self.generate(prompt)
 
-# Wrapper for agent execution
+
 class BasicAgent:
     def __init__(self):
         model = LocalFlanT5Model()
-        tools = [
-            DuckDuckGoSearchTool(),
-            WikipediaSearchTool(),
-        ]
+        tools = [DuckDuckGoSearchTool(), WikipediaSearchTool()]
         self.agent = CodeAgent(
             tools=tools,
             model=model,
-            additional_authorized_imports=["math", "datetime"]
+            additional_authorized_imports=["math", "datetime"],
         )
 
     def __call__(self, question: str) -> str:
-        # FIX: Provide input as a list of message dicts
-        return self.agent.run([{"role": "user", "content": question}]).strip()
+        # smolagents versions differ:
+        # - some expect agent.run("question")
+        # - some accept messages list
+        try:
+            out = self.agent.run(question)
+        except TypeError:
+            out = self.agent.run([{"role": "user", "content": question}])
+        return str(out).strip()
+
+
+def _get_space_repo_id() -> str:
+    # HF Spaces commonly exposes one of these
+    return (
+        os.getenv("SPACE_REPO_NAME")  # often "username/space"
+        or os.getenv("SPACE_ID")
+        or ""
+    ).strip()
+
 
-# Core function to run evaluation
 def run_and_submit_all(profile: gr.OAuthProfile | None) -> tuple[str, pd.DataFrame]:
     if not profile:
         return "Please log in with your Hugging Face account.", pd.DataFrame()
 
-    space_id = os.getenv("SPACE_ID", "your-username/your-space")
-    if space_id == "your-username/your-space":
-        return "Error: Please set the SPACE_ID environment variable.", pd.DataFrame()
+    space_id = _get_space_repo_id()
+    if not space_id or "/" not in space_id:
+        return (
+            "Error: Could not detect Space repo id. "
+            "Make sure you're running inside a HF Space (and not locally).",
+            pd.DataFrame(),
+        )
 
     username = profile.username
-    agent_code_url = f"https://huggingface.co/spaces/{space_id}/tree/main"
+    agent_code_url = f"https://huggingface.co/spaces/{space_id}/blob/main/app.py"
 
     try:
         agent = BasicAgent()
     except Exception as e:
         return f"Error initializing agent: {e}", pd.DataFrame()
 
+    # Fetch questions
     try:
-        response = requests.get(f"{DEFAULT_API_URL}/questions", timeout=15)
-        response.raise_for_status()
-        questions = response.json()
+        resp = requests.get(f"{DEFAULT_API_URL}/questions", timeout=30)
+        resp.raise_for_status()
+        questions = resp.json()
         if not questions:
             return "No questions received from the API.", pd.DataFrame()
     except Exception as e:
@@ -132,23 +163,17 @@ def run_and_submit_all(profile: gr.OAuthProfile | None) -> tuple[str, pd.DataFra
         question = item.get("question")
         if not task_id or not question:
             continue
+
         try:
             submitted_answer = agent(question)
-            answers_payload.append({
-                "task_id": task_id,
-                "submitted_answer": submitted_answer
-            })
-            results_log.append({
-                "Task ID": task_id,
-                "Question": question,
-                "Submitted Answer": submitted_answer
-            })
+            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
+            results_log.append(
+                {"Task ID": task_id, "Question": question, "Submitted Answer": submitted_answer}
+            )
         except Exception as e:
-            results_log.append({
-                "Task ID": task_id,
-                "Question": question,
-                "Submitted Answer": f"AGENT ERROR: {e}"
-            })
+            results_log.append(
+                {"Task ID": task_id, "Question": question, "Submitted Answer": f"AGENT ERROR: {e}"}
+            )
 
     if not answers_payload:
         return "Agent failed to generate any answers.", pd.DataFrame(results_log)
@@ -156,15 +181,16 @@ def run_and_submit_all(profile: gr.OAuthProfile | None) -> tuple[str, pd.DataFra
     submission_data = {
         "username": username,
         "agent_code": agent_code_url,
-        "answers": answers_payload
+        "answers": answers_payload,
     }
 
+    # Submit
     try:
-        res = requests.post(f"{DEFAULT_API_URL}/submit", json=submission_data, timeout=60)
+        res = requests.post(f"{DEFAULT_API_URL}/submit", json=submission_data, timeout=120)
         res.raise_for_status()
         data = res.json()
         summary = (
-            f"✅ Submission Successful!\n"
+            f" Submission Successful!\n"
             f"User: {data.get('username')}\n"
             f"Score: {data.get('score', 'N/A')}% "
             f"({data.get('correct_count', '?')}/{data.get('total_attempted', '?')})\n"
@@ -174,24 +200,21 @@ def run_and_submit_all(profile: gr.OAuthProfile | None) -> tuple[str, pd.DataFra
     except Exception as e:
         return f"Submission Failed: {e}", pd.DataFrame(results_log)
 
-# Gradio UI definition
+
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Agent Evaluation Runner")
+    gr.Markdown("# Agent Evaluation Runner")
     gr.Markdown(
-        "1. Clone this space and edit `app.py` to customize your agent logic.\n"
-        "2. Log in with your Hugging Face account.\n"
-        "3. Press the button to evaluate and submit all answers."
+        "1) Clone this Space and edit `app.py`.\n"
+        "2) Log in with your HF account.\n"
+        "3) Click the button to evaluate and submit."
     )
 
     gr.LoginButton()
     run_button = gr.Button("Run Evaluation & Submit All Answers")
-    status_output = gr.Textbox(label="Submission Status", lines=5, interactive=False)
+    status_output = gr.Textbox(label="Submission Status", lines=6, interactive=False)
     results_table = gr.DataFrame(label="Agent Responses")
 
-    run_button.click(
-        fn=run_and_submit_all,
-        outputs=[status_output, results_table]
-    )
+    run_button.click(fn=run_and_submit_all, outputs=[status_output, results_table])
 
 if __name__ == "__main__":
-    demo.launch(debug=True, share=False)
+    demo.launch()