Construindo aplicações personalizadas com LLM através de RAG (Retrieve Augmented Generation)

Milton Gama Neto
Data Hackers
Published in
7 min readFeb 1, 2024

Os LLMs (Large Language Models) tem apresentado uma eficiência surpreendente em diversas atividades. Isso ocorre porque esses modelos são treinados em enormes bases de dados de diferentes contextos, como informações do Wikipédia, artigos científicos, blogs, códigos, entre outros. Porém para alguns fins específicos, assim como para dados corporativos e/ou privados, o modelo não consegue um bom desempenho, pois não teve acesso a informação. Com isso, um possível resultado são os erros conhecidos como alucinações.

Para aumentar a precisão e confiabilidade dos LLMs e fornecer um modo de expandir sua base de conhecimento, surgiu a técnica Retrieval Augmented Generation, ou simplesmente RAG, proposta por pesquisadores do Facebook AI Research (agora Meta AI). Com essa técnica, é possível construir soluções com dados privados e utilizar a “inteligência” de um LLM sem precisar treiná-lo novamente.

Retrieve Augmented Generation (RAG)

RAG, ou Geração Aumentada de Recuperação em português, é uma técnica que combina a recuperação de informação com um modelo de geração de texto para aprimorar os resultados. Esse método fornece ao modelo informações direcionadas, que podem estar mais atualizadas que o conjunto utilizado em seu treinamento ou serem específicas de um contexto particular.

A estrutura de uma solução de RAG envolve três etapas principais: (1) a identificação de documentos relevantes que refletem o contexto da entrada ou pergunta, (2) a junção deste contexto com um prompt contendo instruções específicas, e finalmente, (3) a criação do texto utilizando um LLM. A imagem a seguir exemplifica cada uma dessas fases:

Estrutura de uma solução RAG. Fonte: Autor

Detalhando um pouco mais cada componente:

- Retrieve

O principal diferencial da técnica de RAG em comparação a utilização de um LLM de maneira convencional, é justamente a etapa de recuperação da informação. Nesse momento, é possível adicionar conteúdo além do conhecimento do LLM.

Essa etapa é fundamental para obter boas respostas, por isso, a busca precisa ser precisa e localizar o conteúdo correto de acordo com a pergunta ou consulta realizada. Ainda que a imagem acima exiba apenas o fluxo percorrido pelo processamento de uma pergunta até chegar em uma resposta, vale salientar que a etapa de retrieve pode ser separada em dois momentos básicos: (1) inserção dos documentos e (2) busca dos documentos relevantes.

Dado a relevância da busca, um componente essencial que elevou a qualidade de buscas textuais e está constantemente presente nas arquiteturas de RAG são os Bancos de Dados Vetoriais, que armazenam e buscam os vetores gerados por meio de embeddings. É importante entender esses dois conceitos também:

Embedding: técnica para representar palavras, textos ou outros objetivos como imagens e vídeos, em uma projeção no formato de vetor de números reais. Essa representação é realizada em um espaço multidimensional, em que os pontos mais próximos são mais semelhantes. Os Embeddings podem ser gerados a partir de Redes Neurais Profundas, que aprendem os padrões semânticos dos dados para realizar a conversão de maneira automática.

Essa representação consegue capturar o valor semântico ou contextual, evitando a dependência de uma correspondência textual exata.

Exemplo de embeddings. Fonte: OpenAI

Banco de Dados Vetorial: é um tipo de banco de dados projetado para salvar informações de vetores multidimensionais. Além do armazenamento, um dos principais benefícios é a forma eficiente de localizar e recuperar informação de acordo com a proximidade ou semelhança entre os vetores. Dessa maneira, são retornadas informações com maior semelhança semântica.

- Augment

Etapa voltada para prompt engineering. O aumento ocorre através da concatenação dos documentos relevantes em um prompt que fornece as instruções para o modelo, indicando o que fazer e que deve considerar as informações listadas assim como a pergunta do usuário.

Com a construção de bons prompts e informações recuperadas de maneira precisa para a pergunta em questão, o modelo receberá um conteúdo que não fez parte do seu treinamento e serão importantes para a resposta final.

- Generate

Com o prompt preparado, basta selecionar o LLM e submetê-lo para gerar a resposta. Dessa maneira, o LLM consegue utilizar suas informações internas (do treinamento) e as informações externas (do prompt) para conseguir gerar uma boa resposta para o que foi solicitado.

Vantagens em utilizar RAG

  • Fornece informações para o LLM que não fez parte do seu treinamento;
  • Governança simples, seja para adicionar, remover ou atualizar o conteúdo;
  • Sem necessidade de retreinar um LLM, ou seja, apresenta um modo mais barato de personalizar a solução;
  • Possibilita a identificação das fontes utilizadas nas respostas, passando transparência para quem utiliza e reduz a chance de alucinação.

Implementando um sistema de RAG

Essa seção apresenta uma implementação de RAG em Python. Foi utilizado o banco de dados vetorial ChromaDB, e os modelos da OpenAI text-embedding-ada-002 e GPT-3.5, para criação de vetores e geração de texto, respectivamente. Entretanto, todos esses componentes podem ser substituídos por outros de mesmo propósito.

Step 0: Upload de documentos

Antes de recuperar as informações, é necessário preparar a base de dados e inserir os documentos em uma collection/index. Além da seleção dos documentos, nesse momento também são construídos os embeddings.

O arquivo prep_docs.py abaixo é responsável por ler os PDFs de uma pasta, transformar e inserir as informações no ChromaDB.

import PyPDF2
from openai import OpenAI
import chromadb
import uuid
import os

CHUNK_SIZE = 1000
OFFSET = 200

openai_client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

chromadb_path = "path/to/save/" # CONFIG YOUR PATH
chroma_client = chromadb.PersistentClient(path=chromadb_path)
collection = chroma_client.create_collection(name="my_collection")

def get_document(document_path):
"""Read a PDF document and return text in string"""
file = open(document_path, 'rb')
reader = PyPDF2.PdfReader(file)

document_text = ""
for i in range(len(reader.pages)):
page = reader.pages[i]
content = page.extract_text()
document_text += content

len(document_text)
return document_text

def split_document(document_text):
"""Split a document in a list of string"""
documents = []
for i in range(0, len(document_text), CHUNK_SIZE):
start = i
end = i + 1000
if start != 0:
start = start - OFFSET
end = end - OFFSET
documents.append(document_text[start:end])
return documents

def get_embedding(text):
"""Transform a text in a vector using embedding model"""
embedding = openai_client.embeddings.create(
input=text,
model="text-embedding-ada-002"
)

return embedding.data[0].embedding

def prepare_documents(documents, document_name):
"""Prepare documents to vector database, generating embedding and metadata"""
embeddings = []
metadatas = []
for i, doc in enumerate(documents):
embeddings.append(get_embedding(doc))
metadatas.append({"source": document_name, "partition" : i})

return embeddings, metadatas

def create_ids(documents):
"""Create a list of IDs for documents"""
return [uuid.uuid4() for _ in documents]

def insert_data(documents, embeddings, metadatas, ids):
"""Insert data in a ChromaDB collection"""
collection.add(
embeddings=embeddings,
documents=documents,
metadatas=metadatas,
ids=ids
)
print(f"Data successfully entered! {len(documents)} Chunks")

def run():
print("Running prep docs...")
data_path = 'data/'

documents = []
embeddings = []
metadatas = []

documents_names = os.listdir(data_path)
documents_names_size = len(documents_names)
for i, document_name in enumerate(documents_names):
print(f"{i+1}/{documents_names_size}: {document_name}")

document = get_document(os.path.join(data_path, document_name))
document_chunks = split_document(document)
document_embeddings, document_metadatas = prepare_documents(document_chunks, document_name)
documents.extend(document_chunks)
embeddings.extend(document_embeddings)
metadatas.extend(document_metadatas)

ids = create_ids(documents)
insert_data(documents, embeddings, metadatas, ids)

if __name__ == "__main__":
run()

As etapas seguintes que implementam o RAG serão divididas no arquivo main.py abaixo. O código é responsável por realizar a consulta no ChromaDB, preparar o prompt e utilizar o LLM.

Step 1: Retrieve

import chromadb
from openai import OpenAI
import os

openai_client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

chromadb_path = "path/to/save/" # CONFIG YOUR PATH
chroma_client = chromadb.PersistentClient(path=chromadb_path)
collection = chroma_client.get_collection("my_collection")

def get_embedding(text):
"""Transform a text in a vector using embedding model"""
embedding = openai_client.embeddings.create(
input=text,
model="text-embedding-ada-002"
)

return embedding.data[0].embedding

def search_document(question):
"""Search for relevant documents in ChromaDB according to the question"""
query_embedding = get_embedding(question)
results = collection.query(
query_embeddings=[query_embedding],
n_results=3
)
return results

def format_search_result(relevant_documents):
"""Auxiliary function to format a list of documents into a string"""
formatted_list = []
for i, doc in enumerate(relevant_documents["documents"][0]):
formatted_list.append("[{}]: {}".format(relevant_documents["metadatas"][0][i]["source"], doc))

documents_str = "\n".join(formatted_list)
return documents_str
relevant_documents = search_document(question)
documents_str = format_search_result(relevant_documents)

Step 2: Augment

prompt = """Você é um assistente de IA que responde as dúvidas dos usuários com bases nos documentos a baixo.
Os documentos abaixo apresentam as fontes atualizadas e devem ser consideradas como verdade.
Cite a fonte quando fornecer a informação.
Documentos:
{documents}
"""

prompt = prompt.format(documents=documents_str)

messages=[
{"role": "system","content": f"{prompt}"},
{"role": "user","content": f"{question}"}
]

Step 3: Generate

chat_completion = openai_client.chat.completions.create(
messages=messages,
model="gpt-3.5-turbo",
max_tokens=512,
temperature=0
)
answer = chat_completion.choices[0].message.content

Código completo

Como evoluir a solução

Como foi mencionado, existem diversas maneiras de adaptar a solução, possibilidades para modificar o banco vetorial ou o LLM utilizado, seja optando por uma solução paga ou por meio de open-source. Cada opção apresentará um tradeoff.

Algumas opções comuns para variar esses componentes na aplicação:

  • Bancos Vetoriais
    - ChromaDB
    - Azure AI Search (Microsoft)
    - Vertex AI Search (Google)
    - Pinecone
    - Milvus
    - Qdrant
  • Modelos (LLMs)
    - GPT-3.5 (OpenAI)
    - GPT-4 (OpenAI)
    - PaLM 2 (Google)
    - LLaMA 2
    - Falcon
    - Mistral 7B

A comparação vetorial é uma excelente maneira de capturar a semântica, porém a comparação textual também tem suas vantagens e por isso combinar as duas, numa abordagem conhecida como busca híbrida, pode extrair o potencial de ambas. Já a escolha do LLM pode ser guiada pelo tamanho do modelo, precificação e em benchmark de performance.

Um tipo comum de aplicação é um chatbot, que mantém o contexto da conversa, isto é, não é executado apenas no formato de Pergunta e Resposta. Com algumas adaptações no código de RAG apresentado, a solução será capaz de utilizar o histórico da conversa. Existem algumas maneiras de fazer isso, por exemplo, utilizando o LLM para sumarizar o histórico ou por meio de janela deslizante enviar as últimas interações.

Uma outra forma de evoluir com o código, é por meio do framework LangChain, que tem o propósito de simplificar a construção de aplicações com LLMs, fornecendo interações e uma abstração de algumas etapas e complexidades. Devido ao intuito do artigo em detalhar as etapas que ocorrem em uma solução de RAG, o framework não foi adotado para conseguir explorar melhor o funcionamento ponta a ponta, porém é fortemente recomendado para escalar e evoluir com as soluções.

Conclusão

RAG é uma excelente maneira de criar aplicações com dados proprietários, além de ter maior domínio em relação ao tipo de informação utilizada. Esse tipo de técnica habilita uma série de aplicações, com um custo significativamente menor em comparação a treinar um modelo especializado.

--

--

Milton Gama Neto
Data Hackers

Cientista de Dados no Nubank e Mestre em Ciência da Computação