Update app.py

app.py

@@ -1,13 +1,11 @@
 import os
-from typing import Optional, Tuple, Dict
+from typing import Optional, Tuple
 import gradio as gr
 from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain_community.llms import HuggingFacePipeline
-from langchain.chains import RetrievalQA
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import tempfile
 import time
@@ -19,6 +17,7 @@ DOCS_DIR = "documents"
 
 class RAGSystem:
     def __init__(self):
+        # Carrega o modelo e o tokenizador
         self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
         self.model = AutoModelForSeq2SeqLM.from_pretrained(
             LLM_MODEL,
@@ -26,21 +25,13 @@ class RAGSystem:
             torch_dtype=torch.float32
         )
         
-        pipe = pipeline(
-            "text2text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
-            max_length=512,
-            temperature=0.3,  # Respostas mais precisas
-            top_p=0.9,        # Diversidade controlada
-            repetition_penalty=1.2  # Evita repetições
-        )
-        
-        self.llm = HuggingFacePipeline(pipeline=pipe)
+        # Configurações de embedding
         self.embeddings = HuggingFaceEmbeddings(
             model_name=EMBEDDING_MODEL,
             model_kwargs={'device': 'cpu'}
         )
+        
+        # Carrega a base de conhecimento
         self.base_db = self.load_base_knowledge()
 
     def load_base_knowledge(self) -> Optional[FAISS]:
@@ -49,6 +40,7 @@ class RAGSystem:
                 os.makedirs(DOCS_DIR)
                 return None
 
+            # Carrega documentos da pasta
             loader = DirectoryLoader(
                 DOCS_DIR,
                 glob="**/*.pdf",
@@ -59,6 +51,7 @@ class RAGSystem:
             if not documents:
                 return None
             
+            # Divide os documentos em trechos menores
             text_splitter = RecursiveCharacterTextSplitter(
                 chunk_size=500,
                 chunk_overlap=100,
@@ -67,6 +60,7 @@ class RAGSystem:
             )
             texts = text_splitter.split_documents(documents)
             
+            # Cria o banco de dados de embeddings
             return FAISS.from_documents(texts, self.embeddings)
             
         except Exception as e:
@@ -75,10 +69,12 @@ class RAGSystem:
 
     def process_pdf(self, file_content: bytes) -> Optional[FAISS]:
         try:
+            # Salva o PDF temporariamente
             with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
                 tmp_file.write(file_content)
                 tmp_path = tmp_file.name
 
+            # Carrega o PDF
             loader = PyPDFLoader(tmp_path)
             documents = loader.load()
             os.unlink(tmp_path)
@@ -86,6 +82,7 @@ class RAGSystem:
             if not documents:
                 return None
             
+            # Divide o PDF em trechos menores
             text_splitter = RecursiveCharacterTextSplitter(
                 chunk_size=500,
                 chunk_overlap=100,
@@ -94,8 +91,10 @@ class RAGSystem:
             )
             texts = text_splitter.split_documents(documents)
             
+            # Cria o banco de dados de embeddings
             db = FAISS.from_documents(texts, self.embeddings)
             
+            # Combina com a base de conhecimento existente, se houver
             if self.base_db is not None:
                 db.merge_from(self.base_db)
             
@@ -105,26 +104,6 @@ class RAGSystem:
             print(f"Erro ao processar PDF: {str(e)}")
             return None
 
-    def format_response(self, raw_response: str, source_type: str, context_found: bool) -> str:
-        """Formata a resposta para um formato padronizado e claro"""
-        if not context_found:
-            return "🔍 Não foram encontradas informações suficientes nos documentos para responder esta pergunta."
-        
-        prefix = ""
-        if source_type == "pdf":
-            prefix = "📄 [Resposta baseada no PDF enviado]\n\n"
-        elif source_type == "base":
-            prefix = "📚 [Resposta baseada na base de documentos]\n\n"
-        elif source_type == "both":
-            prefix = "📚📄 [Resposta baseada em ambas as fontes]\n\n"
-            
-        # Limpa e formata a resposta
-        response = raw_response.strip()
-        if not response:
-            return "🔍 Não foi possível gerar uma resposta adequada com as informações disponíveis."
-            
-        return f"{prefix}{response}"
-
     def generate_response(self, file_obj, query: str, progress=gr.Progress()) -> Tuple[str, str, str]:
         """Retorna (resposta, status, tempo_decorrido)"""
         if not query.strip():
@@ -152,16 +131,14 @@ class RAGSystem:
             
             progress(0.4, desc="Buscando informações relevantes...")
             
-            # Configuração do RAG
+            # Recupera os trechos relevantes
             retriever = db.as_retriever(
                 search_kwargs={
-                    "k": 6,  # Aumenta o número de trechos recuperados
+                    "k": 6,  # Número de trechos recuperados
                     "fetch_k": 10,
                     "score_threshold": 0.5  # Limiar de relevância
                 }
             )
-            
-            # Recupera o contexto
             context_docs = retriever.get_relevant_documents(query)
             
             # Verifica se o contexto é relevante
@@ -173,7 +150,7 @@ class RAGSystem:
             
             progress(0.6, desc="Gerando resposta...")
             
-            # Prompt mais estruturado
+            # Cria o prompt
             prompt = f"""Instruções:
             1. Analise cuidadosamente o contexto fornecido.
             2. Responda à seguinte pergunta em português de forma clara e direta: {query}
@@ -187,20 +164,21 @@ class RAGSystem:
 
             Pergunta: {query}"""
             
-            # Gera resposta
-            result = self.llm(prompt)
-            
-            # Formata a resposta
-            formatted_response = self.format_response(
-                result,
-                source_type,
-                context_found=True
+            # Gera a resposta usando o modelo diretamente
+            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
+            outputs = self.model.generate(
+                inputs["input_ids"],
+                max_length=512,
+                temperature=0.3,
+                top_p=0.9,
+                repetition_penalty=1.2
             )
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             
             elapsed_time = f"{time.time() - start_time:.1f}s"
             progress(1.0, desc="Concluído!")
             
-            return formatted_response, "✅ Sucesso", elapsed_time
+            return response, "✅ Sucesso", elapsed_time
             
         except Exception as e:
             elapsed_time = f"{time.time() - start_time:.1f}s"