Update app.py

app.py

@@ -1,3 +1,4 @@
+from transformers import pipeline
 import os
 import gradio as gr
 import requests
@@ -7,6 +8,7 @@ import pandas as pd
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+HF_TOKEN = os.getenv("HF_TOKEN", None)
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
@@ -19,6 +21,44 @@ class BasicAgent:
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer
 
+class BasicAgent:
+    def __init__(self):
+        # initialize HF inference pipeline once
+        if HF_TOKEN is None:
+            raise ValueError("HF_TOKEN not set in environment")
+        self.generator = pipeline(
+            "text-generation",
+            model="gpt2",               # or your chosen model
+            tokenizer="gpt2",
+            device=-1,                  # CPU; adjust if GPU
+            use_auth_token=HF_TOKEN,
+        )
+        # The GAIA system prompt (no "FINAL ANSWER:" at the end)
+        self.system_prompt = (
+            "You are a general AI assistant. "
+            "I will ask you a question. Answer in as few words as possible—"
+            "either a number, a few words, or a comma-separated list—"
+            "with no extra commentary, prefixes, or units.\n\n"
+        )
+        print("BasicAgent initialized with LLM.")
+
+    def __call__(self, question: str) -> str:
+        # Combine prompt + question
+        prompt = f"{self.system_prompt}Q: {question}\nA:"
+        # Run the model
+        out = self.generator(
+            prompt,
+            max_length=len(prompt.split()) + 20,  # short answer
+            num_return_sequences=1,
+            pad_token_id=self.generator.tokenizer.eos_token_id
+        )
+        # Extract the portion after "A:"
+        text = out[0]["generated_text"]
+        # split on the last "\nA:" or just on "A:" if needed
+        answer = text.split("\nA:")[-1].strip()
+        # final cleanup: no newlines, no trailing punctuation
+        return answer.replace("\n", " ").strip()
+
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,