Update app.py

app.py

@@ -1,64 +1,132 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import PyPDF2
+from sentence_transformers import SentenceTransformer
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+import os
+from typing import List, Tuple
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
+# Inicialização do cliente de inferência e modelo de embeddings
 client = InferenceClient("google/gemma-3-27b-it")
+embedder = SentenceTransformer('all-MiniLM-L6-v2')
 
+# Classe para gerenciar o conhecimento dos PDFs
+class PDFKnowledgeBase:
+    def __init__(self):
+        self.documents = []
+        self.embeddings = None
+        
+    def load_pdfs(self, pdf_directory: str):
+        """Carrega todos os PDFs de um diretório"""
+        self.documents = []
+        for filename in os.listdir(pdf_directory):
+            if filename.endswith('.pdf'):
+                pdf_path = os.path.join(pdf_directory, filename)
+                with open(pdf_path, 'rb') as file:
+                    pdf_reader = PyPDF2.PdfReader(file)
+                    text = ""
+                    for page in pdf_reader.pages:
+                        text += page.extract_text() + "\n"
+                    self.documents.append({
+                        'filename': filename,
+                        'content': text
+                    })
+        
+        # Gera embeddings para todos os documentos
+        contents = [doc['content'] for doc in self.documents]
+        self.embeddings = embedder.encode(contents, convert_to_numpy=True)
+    
+    def get_relevant_context(self, query: str, k: int = 3) -> str:
+        """Recupera os k documentos mais relevantes para a query"""
+        if self.embeddings is None or len(self.documents) == 0:
+            return "Nenhum documento carregado ainda."
+            
+        query_embedding = embedder.encode(query, convert_to_numpy=True)
+        similarities = cosine_similarity([query_embedding], self.embeddings)[0]
+        
+        # Obtém os índices dos k documentos mais similares
+        top_k_indices = np.argsort(similarities)[-k:][::-1]
+        
+        # Constrói o contexto relevante
+        context = ""
+        for idx in top_k_indices:
+            context += f"Documento: {self.documents[idx]['filename']}\n"
+            context += f"Trecho: {self.documents[idx]['content'][:500]}...\n\n"
+        
+        return context
+
+# Inicializa a base de conhecimento
+knowledge_base = PDFKnowledgeBase()
 
 def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
+    message: str,
+    history: List[Tuple[str, str]],
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    pdf_directory: str
 ):
-    messages = [{"role": "system", "content": system_message}]
+    # Carrega os PDFs se ainda não foram carregados
+    if not knowledge_base.documents:
+        knowledge_base.load_pdfs(pdf_directory)
+    
+    # Obtém contexto relevante da base de conhecimento
+    context = knowledge_base.get_relevant_context(message)
+    
+    # Constrói o prompt com o contexto RAG
+    rag_prompt = f"""Você é Grok 3, criado por xAI. Use o seguinte contexto dos documentos para responder à pergunta:
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+{context}
 
-    messages.append({"role": "user", "content": message})
+Pergunta do usuário: {message}
 
-    response = ""
+Responda de forma clara e precisa, utilizando o contexto quando relevante."""
 
-    for message in client.chat_completion(
+    messages = [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": rag_prompt}
+    ]
+    
+    # Adiciona histórico se existir
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
+
+    response = ""
+    for message_chunk in client.chat_completion(
         messages,
         max_tokens=max_tokens,
         stream=True,
         temperature=temperature,
         top_p=top_p,
     ):
-        token = message.choices[0].delta.content
-
+        token = message_chunk.choices[0].delta.content
         response += token
         yield response
 
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+# Interface do Gradio
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(value="Você é um assistente útil que responde com base em documentos PDF.", 
+                  label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, 
+                 label="Top-p (nucleus sampling)"),
+        gr.Textbox(value="./pdfs", label="Diretório dos PDFs"),
     ],
+    title="RAG Chatbot com PDFs",
+    description="Faça perguntas e obtenha respostas baseadas em documentos PDF carregados."
 )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    # Crie um diretório 'pdfs' e coloque seus PDFs lá
+    if not os.path.exists("./pdfs"):
+        os.makedirs("./pdfs")
+        
+    demo.launch()