Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Optional, Tuple
 import gradio as gr
 from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -10,24 +10,22 @@ from langchain.chains import RetrievalQA
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import torch
 import tempfile
+import time
 
 # Configurações
 EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
-LLM_MODEL = "google/flan-t5-large"  # Modelo aberto e sem necessidade de autenticação
+LLM_MODEL = "google/flan-t5-large"
 DOCS_DIR = "documents"
 
 class RAGSystem:
     def __init__(self):
-        # Inicializa o modelo de linguagem
-        print("Carregando modelo de linguagem...")
         self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
         self.model = AutoModelForSeq2SeqLM.from_pretrained(
             LLM_MODEL,
             device_map="auto",
-            torch_dtype=torch.float32  # T5 funciona bem com float32
+            torch_dtype=torch.float32
         )
         
-        # Configura o pipeline
         pipe = pipeline(
             "text2text-generation",
             model=self.model,
@@ -37,28 +35,19 @@ class RAGSystem:
             top_p=0.95
         )
         
-        # Configura o modelo LangChain
         self.llm = HuggingFacePipeline(pipeline=pipe)
-        
-        # Configura embeddings
-        print("Configurando embeddings...")
         self.embeddings = HuggingFaceEmbeddings(
             model_name=EMBEDDING_MODEL,
             model_kwargs={'device': 'cpu'}
         )
-
-        # Carrega a base de conhecimento permanente
         self.base_db = self.load_base_knowledge()
 
     def load_base_knowledge(self) -> Optional[FAISS]:
-        """Carrega a base de conhecimento permanente da pasta de documentos"""
         try:
             if not os.path.exists(DOCS_DIR):
-                print(f"Pasta {DOCS_DIR} não encontrada. Criando...")
                 os.makedirs(DOCS_DIR)
                 return None
 
-            # Carrega todos os PDFs da pasta
             loader = DirectoryLoader(
                 DOCS_DIR,
                 glob="**/*.pdf",
@@ -67,46 +56,35 @@ class RAGSystem:
             documents = loader.load()
             
             if not documents:
-                print("Nenhum documento encontrado na pasta base.")
                 return None
             
-            # Divide o texto em chunks
             text_splitter = RecursiveCharacterTextSplitter(
-                chunk_size=500,  # Chunks menores para o T5
+                chunk_size=500,
                 chunk_overlap=100,
                 length_function=len,
                 separators=["\n\n", "\n", ".", " ", ""]
             )
             texts = text_splitter.split_documents(documents)
             
-            # Cria base de conhecimento
-            print(f"Criando base de conhecimento com {len(texts)} chunks...")
-            db = FAISS.from_documents(texts, self.embeddings)
-            return db
+            return FAISS.from_documents(texts, self.embeddings)
             
         except Exception as e:
             print(f"Erro ao carregar base de conhecimento: {str(e)}")
             return None
 
     def process_pdf(self, file_content: bytes) -> Optional[FAISS]:
-        """Processa o PDF do usuário"""
         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=500,
                 chunk_overlap=100,
@@ -115,10 +93,8 @@ class RAGSystem:
             )
             texts = text_splitter.split_documents(documents)
             
-            # Cria base de conhecimento
             db = FAISS.from_documents(texts, self.embeddings)
             
-            # Se existir uma base permanente, mescla com ela
             if self.base_db is not None:
                 db.merge_from(self.base_db)
             
@@ -128,37 +104,37 @@ class RAGSystem:
             print(f"Erro ao processar PDF: {str(e)}")
             return None
 
-    def generate_response(self, file_obj, query: str) -> str:
-        """Gera resposta para a consulta"""
+    def generate_response(self, file_obj, query: str, progress=gr.Progress()) -> Tuple[str, str, str]:
+        """Retorna (resposta, status, tempo_decorrido)"""
         if not query.strip():
-            return "Por favor, insira uma pergunta."
+            return "Por favor, insira uma pergunta.", "⚠️ Aguardando pergunta", "0s"
         
+        start_time = time.time()
         try:
-            # Se tiver arquivo do usuário, processa e mescla com a base
+            progress(0, desc="Iniciando processamento...")
+            
+            # Processa documento
+            progress(0.2, desc="Processando documento...")
             if file_obj is not None:
                 db = self.process_pdf(file_obj)
                 if db is None:
-                    return "Não foi possível processar o PDF."
-            # Se não tiver arquivo do usuário, usa só a base permanente
+                    return "Não foi possível processar o PDF.", "❌ Erro no processamento", "0s"
             elif self.base_db is not None:
                 db = self.base_db
             else:
-                return "Nenhuma base de conhecimento disponível. Por favor, faça upload de um PDF ou adicione documentos à pasta base."
+                return "Nenhuma base de conhecimento disponível.", "❌ Sem documentos", "0s"
             
-            # Configura o chain RAG
+            progress(0.4, desc="Buscando informações relevantes...")
             qa_chain = RetrievalQA.from_chain_type(
                 llm=self.llm,
                 chain_type="stuff",
                 retriever=db.as_retriever(
-                    search_kwargs={
-                        "k": 4,  # Aumentamos o k para ter mais contexto
-                        "fetch_k": 6
-                    }
+                    search_kwargs={"k": 4, "fetch_k": 6}
                 ),
                 return_source_documents=True
             )
             
-            # Adiciona contexto sobre a fonte da resposta
+            progress(0.6, desc="Gerando resposta...")
             prompt = f"""Baseado nos documentos fornecidos, responda em português à seguinte pergunta:
             {query}
             
@@ -166,60 +142,127 @@ class RAGSystem:
             Se a resposta vier do PDF enviado, indique isso no início.
             Se não encontrar informações suficientes, indique isso claramente."""
             
-            # Gera resposta
             result = qa_chain({"query": prompt})
-            return result["result"]
+            elapsed_time = f"{time.time() - start_time:.1f}s"
+            
+            progress(1.0, desc="Concluído!")
+            return result["result"], "✅ Sucesso", elapsed_time
             
         except Exception as e:
-            return f"Erro ao gerar resposta: {str(e)}"
+            elapsed_time = f"{time.time() - start_time:.1f}s"
+            return f"Erro ao gerar resposta: {str(e)}", "❌ Erro", elapsed_time
 
-# Interface Gradio
 def create_demo():
     rag = RAGSystem()
     
-    with gr.Blocks() as demo:
-        gr.Markdown("# 📚 Sistema RAG de Consulta a Documentos")
-        gr.Markdown(f"""
-        ### Como usar:
-        1. Os documentos da pasta `{DOCS_DIR}` são usados como base de conhecimento
-        2. Você pode fazer upload de PDFs adicionais para consulta
-        3. Digite sua pergunta e aguarde a resposta
-        4. As respostas são baseadas no conteúdo dos documentos
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                file_input = gr.File(
-                    label="Upload do PDF (opcional)",
-                    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("🔍 Buscar Resposta", variant="primary")
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        with gr.Column(elem_id="container"):
+            # Cabeçalho
+            gr.Markdown(
+                """
+                # 🤖 Assistente de Documentos Inteligente
+                
+                Este sistema usa tecnologia RAG (Retrieval-Augmented Generation) para responder perguntas sobre seus documentos.
+                """
+            )
+            
+            # Área principal
+            with gr.Row():
+                # Coluna de entrada
+                with gr.Column():
+                    with gr.Group():
+                        gr.Markdown("### 📄 Documentos")
+                        file_input = gr.File(
+                            label="Upload de PDF (opcional)",
+                            type="binary",
+                            file_types=[".pdf"],
+                            height=100,
+                        )
+                        info = gr.Markdown(
+                            f"""
+                            ℹ️ Além do upload, o sistema também consulta a pasta `{DOCS_DIR}`
+                            """
+                        )
+                    
+                    with gr.Group():
+                        gr.Markdown("### ❓ Sua Pergunta")
+                        query_input = gr.Textbox(
+                            placeholder="Digite sua pergunta aqui...",
+                            lines=3,
+                            max_lines=6,
+                            show_label=False,
+                        )
+                    
+                    with gr.Row():
+                        clear_btn = gr.Button("🗑️ Limpar", variant="secondary")
+                        submit_btn = gr.Button("🔍 Enviar Pergunta", variant="primary")
+                
+                # Coluna de saída
+                with gr.Column():
+                    with gr.Group():
+                        gr.Markdown("### 📝 Resposta")
+                        with gr.Row():
+                            status_output = gr.Textbox(
+                                label="Status",
+                                value="⏳ Aguardando...",
+                                interactive=False,
+                                show_label=False,
+                            )
+                            time_output = gr.Textbox(
+                                label="Tempo",
+                                value="0s",
+                                interactive=False,
+                                show_label=False,
+                            )
+                        
+                        response_output = gr.Textbox(
+                            label="Resposta",
+                            placeholder="A resposta aparecerá aqui...",
+                            interactive=False,
+                            lines=12,
+                            show_label=False,
+                        )
             
-            with gr.Column(scale=1):
-                output = gr.Textbox(
-                    label="Resposta",
-                    lines=10
+            # Exemplos
+            with gr.Accordion("📚 Exemplos de Perguntas", open=False):
+                gr.Examples(
+                    examples=[
+                        [None, "Qual é o tema principal dos documentos?"],
+                        [None, "Pode resumir os pontos principais?"],
+                        [None, "Quais são as principais conclusões?"],
+                        [None, "Explique o contexto deste documento."],
+                    ],
+                    inputs=[file_input, query_input],
                 )
+            
+            # Rodapé
+            gr.Markdown(
+                """
+                ---
+                ### 🔧 Sobre o Sistema
+                * Usa modelo T5 para geração de respostas
+                * Processamento de documentos com tecnologia RAG
+                * Suporte a múltiplos documentos PDF
+                * Respostas baseadas apenas no conteúdo dos documentos
+                """
+            )
         
+        # Eventos
         submit_btn.click(
             fn=rag.generate_response,
             inputs=[file_input, query_input],
-            outputs=output
+            outputs=[response_output, status_output, time_output],
+        )
+        
+        clear_btn.click(
+            lambda: (None, "", "⏳ Aguardando...", "0s"),
+            outputs=[file_input, query_input, status_output, time_output],
         )
         
-        gr.Examples(
-            examples=[
-                [None, "Qual é o tema principal dos documentos?"],
-                [None, "Pode fazer um resumo dos pontos principais?"],
-                [None, "Quais são as principais conclusões?"]
-            ],
-            inputs=[file_input, query_input]
+        # Limpa a resposta quando a pergunta muda
+        query_input.change(
+            lambda: ("", "⏳ Aguardando...", "0s"),
+            outputs=[response_output, status_output, time_output],
         )
     
     return demo