Update app.py

app.py

@@ -21,52 +21,37 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # --- Basic Agent Definition ---
 class GeneralAgent:
     def __init__(self):
-        print("Initializing GPT-2 based QA agent...")
-        self.model = GPT2LMHeadModel.from_pretrained("gpt2")
-        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+        print("Initializing BERT-based QA agent...")
+        # Cargar el modelo BERT preentrenado en SQuAD para tareas de Pregunta y Respuesta
+        self.qa_pipeline = pipeline("question-answering", model="bert-large-uncased-whole-word-masking-finetuned-squad")
     
     def __call__(self, question: str, context: str = None) -> str:
         """
-        Process the question and generate an answer based on the provided context.
+        Procesa la pregunta y devuelve una respuesta basada en el contexto proporcionado.
+        Si no se proporciona contexto, devuelve un mensaje de error.
         """
         if context is None:
-            context = "No context provided."  # Default context if none is given
+            return "FINAL ANSWER: No context provided."
         
-        # Create a clear, structured prompt
+        # Crear un prompt dentro del contexto que estructure la tarea más explícitamente
         prompt = f"""
-You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 
-YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma-separated list of numbers and/or strings. If you are asked for a number, don't use commas to write your number nor use units such as $ or percent sign unless specified otherwise. 
-If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
-If you are asked for a comma-separated list, apply the above rules depending on whether the element to be put in the list is a number or a string.
+You are a general AI assistant. I will ask you a question based on the provided context. 
+Please provide the answer in a clear and concise manner.
 Question: {question}
 Context: {context}
+Answer:
 """
         
-        inputs = self.tokenizer.encode(prompt, return_tensors="pt")
-        
-        # Generate the answer with GPT-2
-        outputs = self.model.generate(inputs, max_length=500, num_return_sequences=1, no_repeat_ngram_size=2, early_stopping=True)
-        
-        # Decode the generated answer
-        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract the relevant part of the answer
-        final_answer = self._extract_final_answer(answer)
-        
-        return f"FINAL ANSWER: {final_answer}"
-
-    def _extract_final_answer(self, answer: str) -> str:
-        """
-        Extract the relevant part of the generated answer.
-        """
-        final_answer_start = "FINAL ANSWER:"
-        start_idx = answer.find(final_answer_start)
-        
-        if start_idx == -1:
-            return "Error processing question."
+        try:
+            # Usar el pipeline para obtener la respuesta de la pregunta con el contexto
+            result = self.qa_pipeline(question=question, context=prompt)
+            answer = result["answer"]
+        except Exception as e:
+            print(f"Error durante QA: {e}")
+            answer = "Error processing question."
         
-        final_answer = answer[start_idx + len(final_answer_start):].strip()
-        return final_answer.strip()
+        # Devuelve la respuesta final con el formato requerido
+        return f"FINAL ANSWER: {answer}"