Update app.py

app.py

@@ -6,6 +6,8 @@ import pandas as pd
 
 import torch
 from transformers import pipeline
+import wikipediaapi
+import textwrap
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -13,33 +15,48 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
+
 class BasicAgent:
     def __init__(self):
-        print("🔄 Loading Zephyr 7B...")
+        print("Loading Wikipedia...")
+        self.wiki = wikipediaapi.Wikipedia('en')
+        print("Loading Zephyr LLM pipeline...")
         self.llm = pipeline(
             "text-generation",
             model="HuggingFaceH4/zephyr-7b-beta",
             tokenizer="HuggingFaceH4/zephyr-7b-beta",
-            max_new_tokens=200,
+            max_new_tokens=300,
             temperature=0,
             device=0 if torch.cuda.is_available() else -1
         )
-        print("✅ Zephyr 7B Loaded.")
+        print("Ready.")
 
     def __call__(self, question: str) -> str:
-        print(f"🤖 Agent received question: {question}")
+        print(f"🧠 Question: {question}")
 
-        prompt = f"<|system|>You are a helpful assistant.<|user|>{question}<|assistant|>"
+        # 1. Fetch page content
+        page = self.wiki.page("Mercedes Sosa")
+        if not page.exists():
+            return "Wikipedia page not found."
 
-        try:
-            output = self.llm(prompt)[0]["generated_text"]
-            # Post-process to remove prompt from output
-            answer = output.split("<|assistant|>")[-1].strip()
-            print(f"✅ Answer: {answer}")
-            return answer
-        except Exception as e:
-            print(f"❌ Error generating response: {e}")
-            return "I don't know"
+        text = page.text
+        chunks = textwrap.wrap(text, width=2000)  # break into ~2k token-like chunks
+
+        best_answer = ""
+        for chunk in chunks:
+            prompt = (
+                "<|system|>You are a precise assistant using Wikipedia.</s>\n"
+                f"<|user|>{question}\n\nHere is relevant context from Wikipedia:\n{chunk}\n<|assistant|>"
+            )
+
+            result = self.llm(prompt)[0]["generated_text"]
+            answer = result.split("<|assistant|>")[-1].strip()
+
+            if any(char.isdigit() for char in answer):  # naive check for answer-like text
+                best_answer = answer
+                break  # stop early if answer found
+
+        return best_answer or "I don't know"
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """