Update app.py

app.py

@@ -231,10 +231,12 @@ class TextProcessor:
 
 class ModelManager:
     """Gerencia o modelo de IA e processamento de perguntas"""
-    
+
     def __init__(self):
         self.model_name = "deepset/roberta-base-squad2"
         self.device = 0 if torch.cuda.is_available() else -1
+        self.max_tokens_answer = 50  # Máximo de tokens na resposta
+        self.max_tokens_context = 300  # Máximo de tokens no contexto exibido
         self.load_model()
 
     def load_model(self):
@@ -243,9 +245,9 @@ class ModelManager:
             self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
             self.model = AutoModelForQuestionAnswering.from_pretrained(self.model_name)
             self.nlp = pipeline('question-answering',
-                              model=self.model,
-                              tokenizer=self.tokenizer,
-                              device=self.device)
+                                 model=self.model,
+                                 tokenizer=self.tokenizer,
+                                 device=self.device)
             logger.info(f"Modelo {self.model_name} carregado com sucesso")
         except Exception as e:
             logger.error(f"Erro ao carregar modelo: {e}")
@@ -254,7 +256,11 @@ class ModelManager:
     def get_answer(self, question: str, context: str) -> Dict:
         """Processa uma única pergunta/contexto"""
         try:
-            return self.nlp(question=question, context=context)
+            return self.nlp(
+                question=question,
+                context=context,
+                max_answer_len=self.max_tokens_answer  # Limitar resposta
+            )
         except Exception as e:
             logger.error(f"Erro ao processar resposta: {e}")
             return {
@@ -278,13 +284,12 @@ class ModelManager:
 
             answers = []
             with ThreadPoolExecutor() as executor:
-                futures = [executor.submit(self.get_answer, question, chunk) 
-                          for chunk in chunks]
+                futures = [executor.submit(self.get_answer, question, chunk) for chunk in chunks]
                 answers = [future.result() for future in futures]
 
             # Filtrar respostas vazias
             answers = [ans for ans in answers if ans['answer'].strip()]
-            
+
             if not answers:
                 return {
                     'answer': "Não foi possível encontrar uma resposta.",
@@ -295,18 +300,25 @@ class ModelManager:
                 }
 
             best_answer = max(answers, key=lambda x: x['score'])
-            
+
+            # Limitar contexto para o máximo de tokens configurado
+            limited_context = " ".join(
+                self.tokenizer.convert_ids_to_tokens(
+                    self.tokenizer.encode(best_answer.get('context', ""), add_special_tokens=False)[:self.max_tokens_context]
+                )
+            )
+
             return {
                 'answer': best_answer['answer'],
                 'score': best_answer['score'],
                 'confidence': f"{best_answer['score']*100:.2f}%",
-                'context': best_answer.get('context', ""),
-                'page_number': None  # TODO: Implementar rastreamento de página
+                'context': limited_context,
+                'page_number': None  # Adapte conforme necessário
             }
         except Exception as e:
-            logger.error(f"Erro ao processar resposta: {e}")
+            logger.error(f"Erro ao obter melhor resposta: {e}")
             return {
-                'answer': "Erro ao processar a pergunta.",
+                'answer': "Erro ao processar o documento.",
                 'score': 0,
                 'confidence': "0%",
                 'context': "",