DHEIVER commited on
Commit
614358b
·
verified ·
1 Parent(s): 79d7438

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +74 -0
app.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
3
+ from sentence_transformers import SentenceTransformer
4
+ import faiss
5
+ import numpy as np
6
+
7
+ # Carregar o modelo RAG
8
+ tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
9
+ retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True)
10
+ model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
11
+
12
+ # Carregar o modelo de embeddings para o retriever
13
+ embedder = SentenceTransformer('paraphrase-MiniLM-L6-v2')
14
+
15
+ # Criar um índice FAISS para o retriever
16
+ dimension = 384 # Dimensão do embedding
17
+ index = faiss.IndexFlatL2(dimension)
18
+
19
+ # Função para adicionar documentos ao índice FAISS
20
+ def add_documents_to_index(documents):
21
+ embeddings = embedder.encode(documents)
22
+ index.add(np.array(embeddings))
23
+
24
+ # Função para recuperar documentos relevantes
25
+ def retrieve_documents(query, k=5):
26
+ query_embedding = embedder.encode([query])
27
+ distances, indices = index.search(np.array(query_embedding), k)
28
+ return [documents[i] for i in indices[0]]
29
+
30
+ # Função para gerar respostas usando RAG
31
+ def rag_answer(query):
32
+ # Recuperar documentos relevantes
33
+ relevant_docs = retrieve_documents(query)
34
+
35
+ # Preparar a entrada para o modelo RAG
36
+ inputs = tokenizer.prepare_seq2seq_batch(query, return_tensors="pt")
37
+
38
+ # Gerar a resposta
39
+ generated = model.generate(input_ids=inputs["input_ids"])
40
+ answer = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
41
+
42
+ return answer, relevant_docs
43
+
44
+ # Interface Gradio
45
+ def gradio_interface(query):
46
+ answer, relevant_docs = rag_answer(query)
47
+ return answer, relevant_docs
48
+
49
+ # Exemplo de documentos (substitua por seus próprios documentos)
50
+ documents = [
51
+ "Aprendizado de máquina é um subcampo da inteligência artificial.",
52
+ "Gradio é uma biblioteca para criar interfaces de usuário para modelos de machine learning.",
53
+ "Transformers são modelos de deep learning usados em NLP.",
54
+ "FAISS é uma biblioteca para busca eficiente de similaridade.",
55
+ "RAG é um modelo que combina recuperação e geração para tarefas de QA."
56
+ ]
57
+
58
+ # Adicionar documentos ao índice FAISS
59
+ add_documents_to_index(documents)
60
+
61
+ # Criar a interface Gradio
62
+ iface = gr.Interface(
63
+ fn=gradio_interface,
64
+ inputs=gr.inputs.Textbox(lines=2, placeholder="Digite sua pergunta aqui..."),
65
+ outputs=[
66
+ gr.outputs.Textbox(label="Resposta"),
67
+ gr.outputs.Textbox(label="Documentos Relevantes")
68
+ ],
69
+ title="Sistema RAG com Gradio",
70
+ description="Um sistema de Retrieval-Augmented Generation (RAG) para responder perguntas."
71
+ )
72
+
73
+ # Iniciar o aplicativo Gradio
74
+ iface.launch()