import gradio as gr
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

# Carregar o modelo RAG
tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True)
model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)

# Carregar o modelo de embeddings para o retriever
embedder = SentenceTransformer('paraphrase-MiniLM-L6-v2')

# Criar um índice FAISS para o retriever
dimension = 384  # Dimensão do embedding
index = faiss.IndexFlatL2(dimension)

# Função para adicionar documentos ao índice FAISS
def add_documents_to_index(documents):
    embeddings = embedder.encode(documents)
    index.add(np.array(embeddings))

# Função para recuperar documentos relevantes
def retrieve_documents(query, k=5):
    query_embedding = embedder.encode([query])
    distances, indices = index.search(np.array(query_embedding), k)
    return [documents[i] for i in indices[0]]

# Função para gerar respostas usando RAG
def rag_answer(query):
    # Recuperar documentos relevantes
    relevant_docs = retrieve_documents(query)
    
    # Preparar a entrada para o modelo RAG
    inputs = tokenizer.prepare_seq2seq_batch(query, return_tensors="pt")
    
    # Gerar a resposta
    generated = model.generate(input_ids=inputs["input_ids"])
    answer = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
    
    return answer, relevant_docs

# Interface Gradio
def gradio_interface(query):
    answer, relevant_docs = rag_answer(query)
    return answer, relevant_docs

# Exemplo de documentos (substitua por seus próprios documentos)
documents = [
    "Aprendizado de máquina é um subcampo da inteligência artificial.",
    "Gradio é uma biblioteca para criar interfaces de usuário para modelos de machine learning.",
    "Transformers são modelos de deep learning usados em NLP.",
    "FAISS é uma biblioteca para busca eficiente de similaridade.",
    "RAG é um modelo que combina recuperação e geração para tarefas de QA."
]

# Adicionar documentos ao índice FAISS
add_documents_to_index(documents)

# Criar a interface Gradio
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Digite sua pergunta aqui..."),
    outputs=[
        gr.outputs.Textbox(label="Resposta"),
        gr.outputs.Textbox(label="Documentos Relevantes")
    ],
    title="Sistema RAG com Gradio",
    description="Um sistema de Retrieval-Augmented Generation (RAG) para responder perguntas."
)

# Iniciar o aplicativo Gradio
iface.launch()