Construindo uma RAG com Gemini e LlamaIndex

Fonte: https://www.llamaindex.ai/blog/llamaindex-gemini-8d7c3b9ea97e

Ultimamente, a personalização dos modelos grandes de linguagem(LLMs) aos contextos específicos tem sido um tópico central. Enquanto os modelos de LLMs oferecem bons resultados, suas aplicações em cenários mais especializados exigem uma especialização adicional. Duas técnicas se destacam nesse processo: Fine-tuning e Retrieval Augmented Generation (RAG):

• Fine-tuning: Consiste em retreinar um modelo pré-treinado com um conjunto de dados específico para uma tarefa. Ela não será abordada nesse artigo.
• RAG: Essa abordagem combina a recuperação de informações relevantes com a geração de texto. Ao invés de depender apenas do conhecimento pré-treinado do modelo, o RAG busca informações em fontes externas, como bases de dados ou documentos, para enriquecer as respostas geradas.

Dada a complexidade e custo de se realizar Fine-Tuning, a RAG tem se tornado uma excelente alternativa para diversas soluções e novos frameworks tornam essa tarefa mais simples. Um desses que vem se tornando popular, pelo menos este ano de 2024 (tudo pode mudar no ano que vem), é o LlamaIndex. Um detalhe bem bacana é que é possível integrar o LlamaIndex com vários modelos de LLM. Nesse tutorial, vou ensinar a montar uma RAG com LlamaIndex e o Gemini, o modelo de LLM do Google. Todo o código está disponível no repositório do Github.

Montando o Código

Estão especificados no notebook todos os pacotes utilizados, inclusive com uso de dotenv. Para o usar modelo Gemini é necessário gerar uma chave de API (já escrevi sobre, recomendo a leitura). Neste tutorial eu faço a leitura de um arquivo científico da minha área de pesquisa em bioinformática. Não publiquei o PDF do artigo, mas as referências estão listadas para acesso. Os passos para montar a RAG estão descritos a seguir:

• Realizar a Leitura do documento: Faz a leitura do PDF do artigo.
• Carregar o modelo Gemini: O código carrega o modelo Gemini, neste caso foi usado o modelo Gemini 1.5 Pro, mas podem ser utilizadas outras versões. No código do Github, apresento uma forma de você verificar todos os disponíveis.
• Criar um modelo de embedding: O modelo de embedding é usado para representar textos como vetores numéricos, neste caso foi usado um modelo da HuggingFace.
• Dividir textos em pedaços: Os textos são divididos em pedaços menores para facilitar o processamento. Recomendo fortemente uma leitura sobre chunks e como definir tamanhos adequados.
• Criar um índice: O índice é criado para armazenar os documentos e suas representações vetoriais através do VectorStoreIndex.
• Criar um engine de consulta: O engine de consulta permite que você faça perguntas ao índice e obtenha respostas relevantes.

from llama_index.core import SimpleDirectoryReader
from llama_index.core import VectorStoreIndex
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings

from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.node_parser import TokenTextSplitter


documents = SimpleDirectoryReader(
    input_files=["pubmed_document.pdf"]
).load_data()


llm  = Gemini(model="models/gemini-1.5-pro", temperature=0.3, top_p=1, top_k=32)
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

text_splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=20)
splitter = TokenTextSplitter(chunk_size=1024, chunk_overlap=20)


# global settings
Settings.llm = llm
Settings.embed_model = embed_model
Settings.text_splitter = text_splitter

index = VectorStoreIndex.from_documents(documents, show_progress=True)

query_engine = index.as_query_engine()

Em diversos tutorias mais antigos eu encontrei o uso de ServiceContext, porém atualmente o LlamaIndex não utiliza este recurso mais e trabalha somente com a Settings, que permite o uso de variáveis globais. Como o documento é pequeno, será bem rápido o processamento, a opção show_progress=True permite a visualização do progresso do processamento. Finalizado o processamento, já é possível fazer as perguntas.

response = query_engine.query(
    "What are the factors that contribute to the development and 
    spread of antibiotic resistance??"
)
print(str(response))

As respostas ficarão restritas ao conteúdo presente no documento ou qualquer outra base de dados que for utilizada. O documento está em inglês, logo é importante que as perguntas sejam em inglês também. Caso a pergunta fuja do contexto, o modelo responde que não é capaz de responder.

Conclusão

Temos aqui um tutorial de como criar uma RAG, lembrando que existem muito recursos que podem ser utilizados e formas de representar os textos que podem ser muito mais poderosas para um determinado propósito. Sempre recomendo consultar as documentações oficiais. De forma geral o objetivo do tutorial é justamente propor algo simples para demonstrar as potencialidades da técnica, além da possibilidade de integrar com vários outros modelos.

#AISprint

Google Cloud credits are provided for this project.

Referências:

Sobre o modelo Gemini: https://blog.google/technology/ai/google-gemini-ai/

Paper da Meta sobre Retrieval-Augmented Generation : https://research.facebook.com/publications/retrieval-augmented-generation-for-knowledge-intensive-nlp-tasks/

Curso de RAG da DeepLearning AI:
https://www.deeplearning.ai/short-courses/building-evaluating-advanced-rag/