Update app.py

app.py

@@ -20,90 +20,74 @@ footer {visibility: hidden}
 # Inicializar o cliente de inferência
 client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
 
+# Variável global para armazenar o retriever
+global_retriever = None
+
 # Função de pré-processamento de texto
 def preprocess_text(text):
     """Pré-processa o texto removendo ruídos e normalizando."""
-    # Remover números de página (ex.: "Página 1", "Page 1 of 10")
     text = re.sub(r'(Página|Page)\s+\d+(?:\s+of\s+\d+)?', '', text, flags=re.IGNORECASE)
-    # Remover múltiplos espaços e quebras de linha
     text = re.sub(r'\s+', ' ', text).strip()
-    # Normalizar texto (remover acentos e converter para minúsculas)
     text = unidecode(text.lower())
     return text
 
-# Configurar o retriever com pré-processamento e indexação avançada
+# Configurar o retriever
 def initialize_retriever(file_objs, persist_directory="chroma_db"):
     """Carrega documentos PDFs, pré-processa e cria um retriever híbrido."""
+    global global_retriever
     if not file_objs:
-        return None, "Nenhum documento carregado."
+        return "Nenhum documento carregado."
     
-    # Carregar e pré-processar documentos
     documents = []
     for file_obj in file_objs:
         loader = PyPDFLoader(file_obj.name)
         raw_docs = loader.load()
         for doc in raw_docs:
             doc.page_content = preprocess_text(doc.page_content)
-            # Adicionar metadados (exemplo: página e origem)
             doc.metadata.update({"source": os.path.basename(file_obj.name)})
         documents.extend(raw_docs)
     
-    # Dividir em pedaços menores
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=2048, chunk_overlap=128)
     splits = text_splitter.split_documents(documents)
     
-    # Criar embeddings e banco de vetores (Chroma)
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     try:
-        # Tentar carregar um banco existente
         vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
-        vectorstore.add_documents(splits)  # Adicionar novos documentos
+        vectorstore.add_documents(splits)
     except:
-        # Criar um novo banco se não existir
         vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings, persist_directory=persist_directory)
     
-    # Configurar retriever semântico
     semantic_retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
-    
-    # Configurar retriever lexical (BM25)
     bm25_retriever = BM25Retriever.from_documents(splits)
     bm25_retriever.k = 2
     
-    # Combinar em um retriever híbrido
-    ensemble_retriever = EnsembleRetriever(
+    global_retriever = EnsembleRetriever(
         retrievers=[semantic_retriever, bm25_retriever],
-        weights=[0.6, 0.4]  # Mais peso para busca semântica
+        weights=[0.6, 0.4]
     )
     
-    return ensemble_retriever, "Documentos processados com sucesso!"
+    return "Documentos processados com sucesso!"
 
 # Formatar o prompt para RAG
 def format_prompt(message, history, retriever=None, system_prompt=None):
     prompt = "<s>"
-    
-    # Adicionar histórico
     for user_prompt, bot_response in history:
         prompt += f"[INST] {user_prompt} [/INST]"
         prompt += f" {bot_response}</s> "
-    
-    # Adicionar instrução do sistema, se fornecida
     if system_prompt:
         prompt += f"[SYS] {system_prompt} [/SYS]"
-    
-    # Adicionar contexto recuperado, se houver retriever
     if retriever:
         docs = retriever.get_relevant_documents(message)
         context = "\n".join([f"[{doc.metadata.get('source', 'Unknown')}, Page {doc.metadata.get('page', 'N/A')}] {doc.page_content}" for doc in docs])
         prompt += f"[CONTEXT] {context} [/CONTEXT]"
-    
-    # Adicionar a mensagem do usuário
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
 # Função de geração com RAG
 def generate(
-    prompt, history, retriever=None, system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0
+    prompt, history, system_prompt=None, temperature=0.2, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0
 ):
+    global global_retriever
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
@@ -118,10 +102,7 @@ def generate(
         seed=42,
     )
 
-    # Formatar o prompt com contexto RAG
-    formatted_prompt = format_prompt(prompt, history, retriever, system_prompt)
-
-    # Gerar resposta em streaming
+    formatted_prompt = format_prompt(prompt, history, global_retriever, system_prompt)
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
 
@@ -132,13 +113,10 @@ def generate(
 # Interface Gradio com RAG
 def create_demo():
     with gr.Blocks(css=css) as demo:
-        retriever_state = gr.State(value=None)
         status = gr.State(value="Nenhum documento carregado")
 
-        # Título
         gr.Markdown("<h1>RAG Chatbot</h1>")
 
-        # Seção de upload de documentos
         with gr.Row():
             with gr.Column(scale=1):
                 gr.Markdown("### Carregar Documentos")
@@ -146,21 +124,18 @@ def create_demo():
                 process_btn = gr.Button("Processar Documentos")
                 status_output = gr.Textbox(label="Status", value="Nenhum documento carregado")
 
-        # Interface de chat
         chat_interface = gr.ChatInterface(
             fn=generate,
             additional_inputs=[
-                gr.State(value=retriever_state),
                 gr.Textbox(label="System Prompt", placeholder="Digite um prompt de sistema (opcional)", value=None)
             ],
             title="",
         )
 
-        # Evento para processar documentos
         process_btn.click(
             fn=initialize_retriever,
             inputs=[file_input],
-            outputs=[retriever_state, status_output]
+            outputs=[status_output]
         )
 
     return demo