Update app.py

app.py

@@ -14,26 +14,25 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
     def __init__(self):
-        print("Initializing T5-based QA agent with custom prompt...")
-        # Load the T5 model
-        self.qa_pipeline = pipeline("text2text-generation", model="t5-large")
+        print("Initializing BERT-based QA agent...")
+        # Load the BERT model fine-tuned on the SQuAD dataset
+        self.qa_pipeline = pipeline("question-answering", model="bert-large-uncased-whole-word-masking-finetuned-squad")
     
-    def __call__(self, question: str) -> str:
+    def __call__(self, question: str, context: str) -> str:
         """
-        Process the question using T5 with a clear and structured prompt, and return an answer.
+        Process the question using the BERT model fine-tuned on SQuAD, and return an answer based on the context.
         """
-        # Create a clearer and more direct prompt
-        prompt = f"Answer the following question: {question}"
-        
         try:
-            # Use the T5 pipeline to generate an answer
-            result = self.qa_pipeline(prompt)
-            answer = result[0]["generated_text"].strip()
+            # Use the QA pipeline to get an answer based on the context and the question
+            result = self.qa_pipeline(question=question, context=context)
+            answer = result["answer"]
         except Exception as e:
             print(f"Error during QA: {e}")
             answer = "Error processing question."
         
         return answer
+
+        
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,