Update app.py

app.py

@@ -4,6 +4,9 @@ import requests
 import pandas as pd
 from transformers import pipeline
 
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+import torch
+
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
@@ -18,26 +21,54 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # --- Basic Agent Definition ---
 class GeneralAgent:
     def __init__(self):
-        print("Initializing general QA agent with improved model...")
-        # Usamos el modelo de RoBERTa entrenado con SQuAD2.0
-        self.qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
+        print("Initializing GPT-2 based QA agent...")
+        # Cargar modelo y tokenizador de GPT-2
+        self.model = GPT2LMHeadModel.from_pretrained("gpt2")
+        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
     
     def __call__(self, question: str, context: str = None) -> str:
         """
-        Process the question and return an answer based on the given context.
-        Uses a prompt template to provide a clear structure for the answer.
+        Procesa la pregunta y genera una respuesta basada en el contexto proporcionado.
+        Usa un prompt específico para guiar la respuesta del modelo GPT-2.
         """
         if context is None:
             return "FINAL ANSWER: No context provided."
         
-        try:
-            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."
+        # Crear el prompt para el modelo GPT-2
+        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 comma to write your number neither 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 of whether the element to be put in the list is a number or a string.
+Question: {question}
+Context: {context}
+"""
+        
+        # Tokenizar el prompt
+        inputs = self.tokenizer.encode(prompt, return_tensors="pt")
+        
+        # Generar la respuesta con GPT-2
+        outputs = self.model.generate(inputs, max_length=500, num_return_sequences=1, no_repeat_ngram_size=2, early_stopping=True)
+        
+        # Decodificar la salida
+        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extraer la respuesta final
+        final_answer = self._extract_final_answer(answer)
+        
+        return f"FINAL ANSWER: {final_answer}"
+
+    def _extract_final_answer(self, answer: str) -> str:
+        """
+        Extrae la parte relevante de la respuesta generada por GPT-2 según el formato solicitado.
+        """
+        # Buscar la sección que comienza con "FINAL ANSWER:"
+        final_answer_start = "FINAL ANSWER:"
+        start_idx = answer.find(final_answer_start)
+        
+        if start_idx == -1:
+            return "Error processing question."
         
-        return f"FINAL ANSWER: {answer}"
+        # Extraer la respuesta que sigue a "FINAL ANSWER:"
+        final_answer = answer[start_idx + len(final_answer_start):].strip()
+        return final_answer.strip()