app.py

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

# 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: str,
    history: List[Tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    pdf_directory: str
):
    # 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:

{context}

Pergunta do usuário: {message}

Responda de forma clara e precisa, utilizando o contexto quando relevante."""

    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_chunk.choices[0].delta.content
        response += token
        yield response

# Interface do Gradio
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        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.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__":
    # Crie um diretório 'pdfs' e coloque seus PDFs lá
    if not os.path.exists("./pdfs"):
        os.makedirs("./pdfs")
        
    demo.launch()