app.py

import os
import gradio as gr
import requests
import pandas as pd
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 our custom model wrapper
        self.model = CustomTransformersModel(model_id="gpt2")
        self.tools = [DuckDuckGoSearchTool()]
        self.agent = CodeAgent(model=self.model, tools=self.tools)

    def __call__(self, question: str) -> str:
        return self.agent.run(question)


# --- Evaluation Logic ---
def run_and_submit_all(profile: gr.OAuthProfile | None):
    space_id = os.getenv("SPACE_ID")

    if profile:
        username = f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # Create the agent
    try:
        agent = SmolAgentWrapper()
    except Exception as e:
        return f"Error initializing agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    # Fetch questions
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
            return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except Exception as e:
        return f"Error fetching questions: {e}", None

    # Run agent
    results_log = []
    answers_payload = []

    for item in questions_data:
        task_id = item.get("task_id")
        question_text = item.get("question")
        if not task_id or question_text is None:
            continue
        try:
            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:
            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})

    if not answers_payload:
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # Submit answers
    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload
    }

    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        final_status = (
            f"Submission Successful!\n"
            f"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    except Exception as e:
        return f"Submission Failed: {e}", pd.DataFrame(results_log)


# --- Gradio Interface ---
with gr.Blocks() as demo:
    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:** Using a custom implementation to handle chat templates properly.
        """
    )

    gr.LoginButton()
    run_button = gr.Button("Run Evaluation & Submit All Answers")
    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

    run_button.click(
        fn=run_and_submit_all,
        outputs=[status_output, results_table]
    )

if __name__ == "__main__":
    print("-" * 60)
    print("Launching SmolAgent Space...")
    print("-" * 60)
    demo.launch(debug=True, share=False)