Update app.py

app.py

@@ -1,152 +1,170 @@
 import os
+from typing import Optional
 import gradio as gr
 from langchain_community.document_loaders import PyPDFLoader
 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 langchain_community.llms import HuggingFaceHub
-from typing import Optional
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 import tempfile
 
 # Configurações
 EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
-LLM_REPO_ID = "google/flan-t5-large"
-
-def create_temporary_file(file_content: bytes) -> str:
-    """Cria um arquivo temporário a partir dos bytes do arquivo."""
-    try:
-        temp_dir = tempfile.mkdtemp()
-        temp_path = os.path.join(temp_dir, "temp.pdf")
-        
-        with open(temp_path, 'wb') as f:
-            f.write(file_content)
-        
-        return temp_path
-    except Exception as e:
-        raise Exception(f"Erro ao criar arquivo temporário: {str(e)}")
+LLM_MODEL = "mistralai/Mistral-7B-v0.1"
 
-def load_and_process_pdf(pdf_path: str) -> Optional[FAISS]:
-    """
-    Carrega e processa o PDF, com tratamento de erros adequado.
-    """
-    try:
-        # Carrega o PDF
-        loader = PyPDFLoader(pdf_path)
-        documents = loader.load()
+class RAGSystem:
+    def __init__(self):
+        # Inicializa o modelo de linguagem
+        self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            LLM_MODEL,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            load_in_8bit=True  # Usa quantização 8-bit para reduzir uso de memória
+        )
         
-        if not documents:
-            return None
-            
-        # Divide o texto em chunks
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=1000,
-            chunk_overlap=200,
-            length_function=len
+        # Configura o pipeline
+        pipe = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            max_length=2048,
+            temperature=0.7,
+            top_p=0.95,
+            repetition_penalty=1.15
         )
-        texts = text_splitter.split_documents(documents)
         
-        # Cria embeddings e armazena no vetor store
-        embeddings = HuggingFaceEmbeddings(
+        # Configura o modelo LangChain
+        self.llm = HuggingFacePipeline(pipeline=pipe)
+        
+        # Configura embeddings
+        self.embeddings = HuggingFaceEmbeddings(
             model_name=EMBEDDING_MODEL,
             model_kwargs={'device': 'cpu'}
         )
-        db = FAISS.from_documents(texts, embeddings)
+
+    def process_pdf(self, file_content: bytes) -> Optional[FAISS]:
+        """Processa o PDF e cria a base de conhecimento"""
+        try:
+            # Cria arquivo temporário
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
+                tmp_file.write(file_content)
+                tmp_path = tmp_file.name
+
+            # Carrega e processa o PDF
+            loader = PyPDFLoader(tmp_path)
+            documents = loader.load()
+            
+            # Remove arquivo temporário
+            os.unlink(tmp_path)
+            
+            if not documents:
+                return None
+            
+            # Divide o texto em chunks
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=1000,
+                chunk_overlap=200,
+                length_function=len,
+                separators=["\n\n", "\n", ".", " ", ""]
+            )
+            texts = text_splitter.split_documents(documents)
+            
+            # Cria base de conhecimento
+            db = FAISS.from_documents(texts, self.embeddings)
+            return db
+            
+        except Exception as e:
+            print(f"Erro ao processar PDF: {str(e)}")
+            return None
+
+    def generate_response(self, file_obj, query: str) -> str:
+        """Gera resposta para a consulta"""
+        if file_obj is None:
+            return "Por favor, faça upload de um arquivo PDF."
         
-        return db
+        if not query.strip():
+            return "Por favor, insira uma pergunta."
         
-    except Exception as e:
-        print(f"Erro ao processar o PDF: {str(e)}")
-        return None
+        try:
+            # Processa o PDF
+            db = self.process_pdf(file_obj)
+            if db is None:
+                return "Não foi possível processar o PDF."
+            
+            # Configura o chain RAG
+            qa_chain = RetrievalQA.from_chain_type(
+                llm=self.llm,
+                chain_type="stuff",
+                retriever=db.as_retriever(
+                    search_kwargs={
+                        "k": 3,
+                        "fetch_k": 5
+                    }
+                ),
+                return_source_documents=True
+            )
+            
+            # Gera resposta
+            result = qa_chain({"query": query})
+            return result["result"]
+            
+        except Exception as e:
+            return f"Erro ao gerar resposta: {str(e)}"
 
-def generate_response(file_obj, query: str) -> str:
-    """
-    Gera resposta para a consulta usando RAG, com tratamento de erros.
-    """
-    if file_obj is None:
-        return "Erro: Nenhum arquivo PDF foi carregado."
-    
-    if not query.strip():
-        return "Erro: Por favor, insira uma pergunta."
+# Interface Gradio
+def create_demo():
+    rag = RAGSystem()
     
-    try:
-        # Cria arquivo temporário e processa o PDF
-        temp_path = create_temporary_file(file_obj)
-        db = load_and_process_pdf(temp_path)
+    with gr.Blocks() as demo:
+        gr.Markdown("# 📚 Sistema RAG com Mistral-7B")
+        gr.Markdown("""
+        ### Instruções:
+        1. Faça upload de um arquivo PDF
+        2. Digite sua pergunta sobre o conteúdo
+        3. Aguarde a resposta gerada pelo modelo
+        """)
         
-        if db is None:
-            return "Erro: Não foi possível processar o PDF."
+        with gr.Row():
+            with gr.Column(scale=1):
+                file_input = gr.File(
+                    label="Upload do PDF",
+                    type="binary",
+                    file_types=[".pdf"]
+                )
+                query_input = gr.Textbox(
+                    label="Sua Pergunta",
+                    placeholder="Digite sua pergunta sobre o documento...",
+                    lines=3
+                )
+                submit_btn = gr.Button("🔍 Pesquisar", variant="primary")
+            
+            with gr.Column(scale=1):
+                output = gr.Textbox(
+                    label="Resposta",
+                    lines=10
+                )
         
-        # Configura o modelo de linguagem
-        llm = HuggingFaceHub(
-            repo_id=LLM_REPO_ID,
-            huggingfacehub_api_token=os.environ.get("HUGGINGFACE_API_TOKEN"),
-            model_kwargs={
-                "temperature": 0.7,
-                "max_length": 512,
-                "top_p": 0.95
-            }
+        submit_btn.click(
+            fn=rag.generate_response,
+            inputs=[file_input, query_input],
+            outputs=output
         )
         
-        # Cria a cadeia de RAG
-        qa_chain = RetrievalQA.from_chain_type(
-            llm=llm,
-            chain_type="stuff",
-            retriever=db.as_retriever(search_kwargs={"k": 3}),
-            return_source_documents=True,
-            verbose=True
+        gr.Examples(
+            examples=[
+                [None, "Qual é o tema principal deste documento?"],
+                [None, "Pode fazer um resumo dos pontos principais?"],
+                [None, "Quais são as principais conclusões?"]
+            ],
+            inputs=[file_input, query_input]
         )
-        
-        # Executa a consulta
-        result = qa_chain({"query": query})
-        
-        # Limpa arquivos temporários
-        os.remove(temp_path)
-        os.rmdir(os.path.dirname(temp_path))
-        
-        return result["result"]
-        
-    except Exception as e:
-        return f"Erro ao gerar resposta: {str(e)}"
-
-# Interface Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# Sistema de RAG com LangChain")
-    gr.Markdown("Faça upload de um PDF e faça perguntas sobre o conteúdo.")
-    
-    with gr.Row():
-        with gr.Column():
-            file_input = gr.File(
-                label="Upload PDF",
-                type="binary",
-                file_types=[".pdf"]
-            )
-            query_input = gr.Textbox(
-                label="Sua Pergunta",
-                placeholder="Digite sua pergunta aqui...",
-                lines=3
-            )
-            submit_btn = gr.Button("Enviar Pergunta")
-        
-        with gr.Column():
-            output = gr.Textbox(
-                label="Resposta Gerada",
-                lines=10
-            )
-    
-    submit_btn.click(
-        fn=generate_response,
-        inputs=[file_input, query_input],
-        outputs=output
-    )
     
-    gr.Examples(
-        examples=[
-            [None, "Qual é o principal tema deste documento?"],
-            [None, "Pode resumir os pontos principais?"]
-        ],
-        inputs=[file_input, query_input]
-    )
+    return demo
 
 if __name__ == "__main__":
+    demo = create_demo()
     demo.launch()