Update app.py

app.py

@@ -1,147 +1,54 @@
-import os
 import gradio as gr
-import requests
-import re
-import pandas as pd
-from transformers import pipeline, set_seed
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from gaia import run_gaia_evaluation
 
-# ---------------- CONSTANTS ----------------
-DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+MODEL_NAME = "google/flan-t5-base"  # use base for stability
 
-# ---------------- BASIC AGENT ----------------
-class BasicAgent:
-    def __init__(self):
-        print("BasicAgent initialized")
-        set_seed(42)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 
-        self.generator = pipeline(
-            "text2text-generation",
-            model="google/flan-t5-large",
-            max_new_tokens=64,
-            temperature=0.0,
-            do_sample=False
-        )
-
-    def __call__(self, question: str) -> str:
-        prompt = (
-            "Answer the question EXACTLY as expected.\n"
-            "Give the final answer only.\n"
-            "Do not explain.\n\n"
-            f"Question:\n{question}\n\n"
-            "Answer:"
-        )
-
-        result = self.generator(prompt)[0]["generated_text"]
-
-        # ✅ MINIMAL CLEANUP (CRITICAL)
-        answer = result.strip()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-        # Remove only obvious prefixes
-        answer = re.sub(r"(?i)^(answer:|the answer is)\s*", "", answer)
 
-        # Keep full line, keep units, keep words
-        answer = answer.split("\n")[0].strip()
-
-        print(f"\nQ: {question}\nA: {answer}\n")
-        return answer
-
-
-# ---------------- RUN + SUBMIT ----------------
-def run_and_submit_all(profile: gr.OAuthProfile | None):
-    space_id = os.getenv("SPACE_ID")
-
-    if not profile:
-        return "Please login with Hugging Face.", pd.DataFrame()
-
-    username = profile.username
-    agent = BasicAgent()
-
-    agent_code = (
-        f"https://huggingface.co/spaces/{space_id}/tree/main"
-        if space_id else "N/A"
-    )
+def answer_question(question: str) -> str:
+    prompt = f"Answer the following question concisely:\n{question}"
 
-    # Fetch questions
-    try:
-        questions = requests.get(
-            f"{DEFAULT_API_URL}/questions", timeout=15
-        ).json()
-    except Exception as e:
-        return f"Failed to fetch questions: {e}", pd.DataFrame()
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(device)
 
-    answers_payload = []
-    results_log = []
-
-    for q in questions:
-        task_id = q["task_id"]
-        question_text = q["question"]
-
-        try:
-            answer = agent(question_text)
-        except Exception as e:
-            answer = "ERROR"
-
-        answers_payload.append({
-            "task_id": task_id,
-            "submitted_answer": answer
-        })
-
-        results_log.append({
-            "Task ID": task_id,
-            "Question": question_text,
-            "Submitted Answer": answer
-        })
-
-    submission_data = {
-        "username": username,
-        "agent_code": agent_code,
-        "answers": answers_payload
-    }
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=64,
+            do_sample=False
+        )
 
-    # Submit answers
-    try:
-        response = requests.post(
-            f"{DEFAULT_API_URL}/submit",
-            json=submission_data,
-            timeout=60
-        ).json()
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return answer.strip()
 
-        status = (
-            f"Submission Successful!\n"
-            f"User: {response.get('username')}\n"
-            f"Overall Score: {response.get('score')}% "
-            f"({response.get('correct_count')}/"
-            f"{response.get('total_attempted')} correct)\n"
-            f"Message: {response.get('message')}"
-        )
-    except Exception as e:
-        status = f"Submission failed: {e}"
 
-    return status, pd.DataFrame(results_log)
+def run_evaluation():
+    """
+    IMPORTANT:
+    - This function MUST return
+    - Must NOT print
+    - Must NOT loop forever
+    """
+    results = run_gaia_evaluation(answer_question)
+    return results
 
 
-# ---------------- GRADIO UI ----------------
 with gr.Blocks() as demo:
     gr.Markdown("# GAIA Unit 4 – Basic Agent Runner")
 
-    gr.LoginButton()
     run_btn = gr.Button("Run Evaluation & Submit")
-
-    status_box = gr.Textbox(
-        label="Submission Result",
-        lines=5,
-        interactive=False
-    )
-
-    results_table = gr.DataFrame(
-        label="Questions and Answers",
-        wrap=True
-    )
+    output = gr.JSON(label="Submission Result")
 
     run_btn.click(
-        fn=run_and_submit_all,
-        outputs=[status_box, results_table]
+        fn=run_evaluation,
+        inputs=[],
+        outputs=output
     )
 
-if __name__ == "__main__":
-    demo.launch(debug=True)
+demo.launch()