Construyendo RAG con LlamaIndex y Mistral

En este notebook, exploraremos cómo construir un sistema de Generación Aumentada por Recuperación (RAG) utilizando la biblioteca LlamaIndex y el modelo de embeddings MistralAI LLM and Embedding Model. Un sistema RAG es fundamental para que los modelos de lenguaje grandes (LLMs) puedan acceder a información externa y generar respuestas más precisas y contextualizadas, superando las limitaciones de su conocimiento pre-entrenado. Además, veremos cómo utilizar un índice ya creado como un recuperador eficiente.

Objetivos de la lección

Al finalizar esta lección, habrás aprendido a:

1. Construir un pipeline RAG básico desde cero.
2. Configurar y utilizar un índice como un recuperador de información.
3. Integrar modelos de lenguaje y embeddings para mejorar la calidad de las respuestas.

Aquí te presentamos los temas que cubriremos:

1. Pipeline RAG básico.
2. Uso del índice como recuperador.

!pip install llama-index 
!pip install llama-index-embeddings-mistralai
!pip install llama-index-llms-mistralai

Configurar Claves API

Antes de comenzar, es crucial configurar tus claves API para acceder a los servicios de Mistral. Asegúrate de tener tu clave API de Mistral disponible y de establecerla como una variable de entorno o directamente en el código, como se muestra a continuación. Esto permitirá que LlamaIndex se comunique con los modelos de Mistral para generar embeddings y respuestas.

import os
os.environ['MISTRAL_API_KEY'] = '<YOUR MISTRALAI API KEY>'

Pipeline RAG básico

Un pipeline RAG básico es una secuencia de pasos que permite a un LLM recuperar información relevante de una base de conocimiento externa antes de generar una respuesta. Esto mejora significativamente la precisión y la relevancia de las respuestas del modelo. A continuación, se detallan los pasos involucrados en la construcción de un pipeline RAG básico.

1. Configurar el LLM y el modelo de embeddings.
2. Descargar los datos.
3. Cargar los datos.
4. Crear nodos a partir de los datos.
5. Crear un índice a partir de los nodos.
6. Crear un motor de consulta (Query Engine).
7. Realizar consultas.

El motor de consulta (Query Engine) combina los módulos Retrieval y Response Synthesis para generar una respuesta a la consulta dada. El recuperador se encarga de encontrar la información relevante, mientras que el sintetizador de respuesta formula la respuesta final basándose en esa información y la consulta del usuario.

Configurar LLM y Modelo de Embeddings

El primer paso es configurar el modelo de lenguaje grande (LLM) que utilizaremos para generar las respuestas y el modelo de embeddings que transformará el texto en representaciones numéricas. Estas representaciones son esenciales para la búsqueda de similitud en el proceso de recuperación.

from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding

Este fragmento de código importa las clases necesarias de LlamaIndex para interactuar con los modelos de Mistral. Se importan MistralAI para el LLM y MistralAIEmbedding para el modelo de embeddings.

llm = MistralAI(model='mistral-large')
embed_model = MistralAIEmbedding()

Aquí se inicializan el LLM y el modelo de embeddings de Mistral. Es importante especificar el modelo de embeddings que se utilizará, en este caso, mistral-embed. La configuración del LLM se deja por defecto, pero podrías especificar parámetros adicionales si fuera necesario.

from llama_index.core import Settings
Settings.llm = llm
Settings.embed_model = embed_model

Finalmente, se establece la configuración global de LlamaIndex para usar los modelos de Mistral que acabamos de inicializar. Esto asegura que cualquier componente de LlamaIndex que necesite un LLM o un modelo de embeddings utilizará estas instancias.

Descargar Datos

Para esta demostración, utilizaremos Uber 2021 10K SEC Filings como nuestra fuente de datos. Este conjunto de datos nos permitirá simular un escenario real donde el LLM necesita acceder a información específica para responder preguntas.

!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O './uber_2021.pdf'

Este comando de shell descarga el archivo paul_graham_essay.txt desde una URL específica. Este archivo contiene el ensayo de Paul Graham, que servirá como nuestra base de conocimiento para el sistema RAG.

--2024-03-29 13:12:05--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.109.133, 185.199.108.133, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 1880483 (1.8M) [application/octet-stream]
Saving to: ‘./uber_2021.pdf’

./uber_2021.pdf     100%[===================>]   1.79M  --.-KB/s    in 0.1s    

2024-03-29 13:12:05 (13.7 MB/s) - ‘./uber_2021.pdf’ saved [1880483/1880483]

Este comando muestra el contenido de las primeras 1000 líneas del archivo descargado. Esto es útil para verificar que el archivo se ha descargado correctamente y para tener una idea del tipo de datos con los que estamos trabajando.

Cargar Datos

Una vez que los datos han sido descargados, el siguiente paso es cargarlos en un formato que LlamaIndex pueda procesar. LlamaIndex utiliza el concepto de "Documentos" para representar los datos de entrada.

from llama_index.core import SimpleDirectoryReader

Aquí se importa la clase SimpleDirectoryReader de LlamaIndex, que es una herramienta conveniente para cargar documentos desde un directorio local. Esta clase simplifica el proceso de lectura de archivos de texto.

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

Esta línea de código utiliza SimpleDirectoryReader para cargar todos los archivos del directorio actual (./). El resultado es una lista de objetos Document, donde cada objeto representa un archivo de texto cargado. En nuestro caso, cargará el ensayo de Paul Graham.

Crear Nodos

Los "nodos" son unidades más pequeñas y manejables de información derivadas de los documentos originales. Dividir los documentos en nodos permite una recuperación más granular y eficiente. LlamaIndex puede crear nodos automáticamente, lo que facilita la gestión de grandes volúmenes de texto.

from llama_index.core.node_parser import TokenTextSplitter

Se importa la clase SentenceSplitter, que es un transformador de nodos. Este transformador se encarga de dividir el texto de los documentos en nodos más pequeños, generalmente basados en oraciones o párrafos, lo que mejora la precisión de la recuperación.

splitter = TokenTextSplitter(
    chunk_size=512,
    chunk_overlap=0,
)

nodes = splitter.get_nodes_from_documents(documents)

Aquí se inicializa el SentenceSplitter y se utiliza para transformar los documentos cargados en una lista de nodos. El parámetro chunk_size define el tamaño máximo de cada nodo en tokens, y chunk_overlap especifica cuántos tokens se superpondrán entre nodos adyacentes, lo que ayuda a mantener el contexto.

Crear Índice

Un índice es una estructura de datos que permite una búsqueda eficiente de los nodos relevantes. LlamaIndex ofrece varios tipos de índices, pero el VectorStoreIndex es comúnmente utilizado para RAG, ya que almacena los embeddings de los nodos y permite búsquedas de similitud vectorial.

from llama_index.core import VectorStoreIndex

Se importa la clase VectorStoreIndex, que es el tipo de índice que utilizaremos para almacenar y buscar nuestros nodos. Este índice es ideal para escenarios RAG porque permite la búsqueda de similitud basada en embeddings.

index = VectorStoreIndex(nodes)

Esta línea crea el VectorStoreIndex a partir de los nodos que generamos previamente. LlamaIndex se encarga de generar los embeddings para cada nodo utilizando el modelo de embeddings configurado y de almacenarlos en el índice, listos para ser consultados.

Crear Motor de Consulta

El motor de consulta (Query Engine) es el componente que orquesta el proceso de RAG. Recibe una consulta del usuario, utiliza el índice para recuperar los nodos más relevantes y luego pasa esos nodos al LLM para generar una respuesta coherente.

query_engine = index.as_query_engine(similarity_top_k=2)

Aquí se crea el motor de consulta a partir del índice. El motor de consulta está listo para recibir preguntas y utilizar el índice para encontrar la información necesaria para responderlas. Este es el punto de entrada principal para interactuar con el sistema RAG.

Realizar Consultas

Una vez que el motor de consulta está configurado, podemos empezar a hacer preguntas. El sistema RAG se encargará de todo el proceso: recuperar, sintetizar y responder.

response = query_engine.query("What is the revenue of Uber in 2021?")

Esta línea envía una consulta al motor de consulta. El motor de consulta procesará esta pregunta, buscará en el índice los nodos más relevantes y utilizará el LLM para generar una respuesta basada en la información recuperada.

print(response)

Este comando imprime la respuesta generada por el motor de consulta. Deberías ver una respuesta que utiliza la información del ensayo de Paul Graham para responder a la pregunta sobre su experiencia en Y Combinator.

The total revenue for Uber in 2021 was $17,455 million. This includes revenue from various offerings such as Mobility, Delivery, Freight, and All Other revenue streams. The Mobility revenue was $6,953 million, Delivery revenue was $8,362 million, Freight revenue was $2,132 million, and All Other revenue was $8 million.

Aquí se realiza una segunda consulta, esta vez sobre la opinión de Paul Graham acerca de las startups. Esto demuestra la capacidad del sistema RAG para responder a diferentes preguntas utilizando la misma base de conocimiento.

Índice como Recuperador

Además de ser una estructura de almacenamiento, podemos utilizar un índice ya creado como un Retriever. Un recuperador es un componente clave en el pipeline RAG que se encarga de encontrar los "chunks" o nodos más relevantes de tu base de conocimiento para una consulta de usuario dada. Al usar el índice directamente como recuperador, podemos tener un control más granular sobre el proceso de recuperación.

Crear Recuperador

Para utilizar el índice como recuperador, simplemente necesitamos obtener una instancia de recuperador directamente del índice.

retriever = index.as_retriever(similarity_top_k = 2)

Esta línea de código crea una instancia de recuperador directamente desde el índice que ya hemos construido. El recuperador está configurado para utilizar el índice subyacente para realizar búsquedas de similitud y devolver los nodos más relevantes.

Recuperar Nodos Relevantes para una Consulta

Una vez que tenemos el recuperador, podemos usarlo para obtener los nodos más relevantes para cualquier consulta, antes de pasarlos a un LLM para la generación de la respuesta.

retrieved_nodes = retriever.retrieve("What is the revenue of Uber in 2021?")

Aquí se utiliza el recuperador para obtener los nodos más relevantes para la consulta "What did Paul Graham do at Y Combinator?". El recuperador buscará en el índice y devolverá una lista de nodos que son semánticamente similares a la consulta.

from llama_index.core.response.notebook_utils import display_source_node

for node in retrieved_nodes:
    display_source_node(node, source_length=1000)

Este bucle itera sobre los nodos recuperados y los imprime. Esto te permite ver exactamente qué fragmentos de información el recuperador consideró relevantes para la consulta, lo cual es útil para depuración y comprensión.

<IPython.core.display.Markdown object>

Se realiza una segunda recuperación con una consulta diferente para demostrar la flexibilidad del recuperador. Esto muestra cómo el recuperador puede adaptarse a diferentes preguntas y encontrar la información pertinente.

<IPython.core.display.Markdown object>

Similar al paso anterior, este bucle imprime los nodos recuperados para la segunda consulta. Observa cómo los nodos recuperados cambian en función de la pregunta, lo que valida la efectividad del recuperador.

Resumen

• Hemos construido un pipeline RAG básico utilizando LlamaIndex y los modelos de Mistral para mejorar la capacidad de respuesta de los LLMs.
• Aprendimos a configurar el LLM y el modelo de embeddings, cargar datos, crear nodos y construir un índice vectorial.
• Creamos un motor de consulta para interactuar con el sistema RAG y realizamos consultas para obtener respuestas contextualizadas.
• Exploramos cómo utilizar un índice existente como un recuperador para obtener nodos relevantes de manera más controlada.
• Este enfoque permite a los LLMs acceder a información externa, lo que resulta en respuestas más precisas y actualizadas.