RAG con Mistral

En este documento probamos la técnica de RAG con las APIs de Mistral. Queremos ver cómo cambian las respuestas de un modelo (un LLM) cuando a las preguntas que le hacemos les añadimos contexto con la información generada con RAG. Estas pruebas también nos permiten comprobar la facilidad de utilizar las APIs de Mistral. Comprobamos que RAG es un método efectivo para que un modelo pueda incorporar información actualizada. Para nuestras pruebas utilizamos una versión modificada de la Odisea de Homero.

¿Qué es el RAG?

El Retrieval-Augmented Generation (RAG) es una técnica avanzada en inteligencia artificial que combina dos enfoques: la generación de texto y la recuperación de información.

Cuando enviamos una pregunta a un modelo (un LLM), en lugar usar solo el modelo para generar la respuesta, con RAG previamente podemos buscar información relevante en documento para mejorar la precisión y relevancia del contenido que genera el modelo.

RAG es útil para incorporar información que no ha visto el modelo en su fase de entreno; al condicionar los resultados del modelo a la información que le damos, lo estamos sesgando hacia esa dirección. RAG es una forma para que un modelo tenga en cuenta información adicional sin tener que modificar el propio modelo.

¿Cómo funciona el RAG?

El RAG se divide en tres partes principales:

  1. Representación Vectorial: Antes de enviar la información de un documento al modelo, lo convertimos en vectores numéricos mediante técnicas de embeddings. Estas representaciones vectoriales se almacenan en una base de datos vectorial, facilitando la posterior búsqueda de información relevante.
  2. Recuperación (Retrieval): Dada la pregunta que vamos a hacer al modelo, se utiliza un módulo que busca en la base de datos vectorial y obtiene los datos más relevantes. Este módulo aplica técnicas matemáticas como la búsqueda por similitud sobre los embeddings que hemos comentado en el paso previo. Básicamente, cuando haces una pregunta, este módulo se encarga de buscar en la base de datos vectorial la información que más se relaciona con tu pregunta.
  3. Generación (Generation): Una vez que se ha recuperado la información relevante, se utiliza un modelo generativo (como GPT4 o Mistral7b) que toma estos datos y genera una respuesta más precisa y detallada. Esto quiere decir que el modelo no solo inventa la respuesta, sino que la construye utilizando la información que ha encontrado previamente en los documentos, lo que mejora mucho la calidad de la respuesta.

Ventajas del RAG

La principal ventaja de usar RAG es que combina lo mejor de dos mundos: la capacidad de generar texto coherente y natural de los modelos generativos, y la precisión de la información actualizada y relevante de los módulos de recuperación. Esto lo hace ideal para aplicaciones donde es crucial tener información exacta y al día, como en asistentes virtuales, chatbots informativos, y más.

Implementación de RAG en Mistral

Vamos a ver un ejemplo sencillo de cómo implementar RAG usando Mistral, prestando especial atención a la división de datos en fragmentos (chunking y splitting) y la generación de embeddings para la recuperación de información.

Los fragmentos “chunks” son las unidades de información que son buscadas y enviadas al LLM cuando le hacemos una pregunta. En concreto, en este ejemplo vamos a implementar un caso de uso de RAG condicionando los resultados del modelo con “chunks” de una versión modificada de la Odisea de Homero.

¿Por qué utilizamos una versión modificada? Un LLM como los de Mistral ha sido entrenado con información online, lo que incluye las versiones libres de la Odisea de Homero publicadas en internet. Esto significa que el modelo ya ha “visto” la Odisea en su fase de entreno y conoce su contenido. Con RAG podemos incorporar una versión modificada. Si enviamos esta versión modificada al modelo, éste nos devolverá respuestas distintas.

Nuestra versión “actualizada” de la Odisea incluye el siguiente párrafo al final: “And Zeus pitying Achilles decided to give him back his life. Achilles left the underworld and spent the rest of his life as a shepperd in the hills of Ithaca”.

Contrastaremos las respuestas del modelo con y sin RAG a la pregunta: “What was Achilles’ ending?”.

Paso 1: Preparación y división del texto

En un sistema RAG, es crucial dividir el documento en fragmentos más pequeños (chunks) para identificar y recuperar la información más relevante de manera más efectiva. En nuestro caso, dividimos el texto en fragmentos de 2048 caracteres, obteniendo 295 fragmentos.

import requests

# Obtenemos el texto desde un repositorio
response = requests.get(--ruta donde alamcenamos la versión modificada de la Odisea--)
text = response.text

# Dividimos el texto en fragmentos de tamaño fijo
chunk_size = 2048
chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
Output:
len(chunks)
295

El tamaño de los fragmentos debe ajustarse según cada caso específico, ya que fragmentos más pequeños pueden mejorar la precisión del retrieval al evitar texto irrelevante, pero también aumentan el tiempo y los recursos de procesamiento. Por lo tanto, se debe encontrar un punto medio entre rendimiento y coste para lograr un resultado óptimo en RAG.

Paso 2: Generación de embeddings

Creamos representaciones numéricas (embeddings) para cada fragmento, utilizando la API de Mistral.

from mistralai.client import MistralClient
import numpy as np

client = MistralClient(api_key='your_api_key')

# Función para obtener embeddings
def get_text_embedding(input):
    embeddings_batch_response = client.embeddings(
        model="mistral-embed",
        input=input
    )
    return embeddings_batch_response.data[0].embedding

# Generamos embeddings para cada fragmento
text_embeddings = np.array([get_text_embedding(chunk) for chunk in chunks])

Vemos que los embeddings generados no dejan de ser un conjunto de vectores, y los fragmentos más cercanos en significado tendrán representaciones vectoriales más próximas:

Output:
text_embeddings
array([[-3.95202637e-02,  7.75756836e-02, -8.82148743e-05, ...,
        -1.26342773e-02, -2.12402344e-02, -2.50816345e-03],
       [-3.19213867e-02,  7.21435547e-02,  2.99835205e-02, ...,
        -1.08413696e-02, -1.19628906e-02, -7.66372681e-03],
       [-5.89599609e-02,  6.12487793e-02,  1.26419067e-02, ...,
        -2.25372314e-02,  4.67681885e-03, -6.26754761e-03],

Paso 3: Almacenamiento en una base de datos vectorial

Almacenamos los embeddings en una base de datos vectorial como Faiss para facilitar la búsqueda.

import faiss

# Creamos un índice en Faiss y añadimos los embeddings
d = text_embeddings.shape[1]
index = faiss.IndexFlatL2(d)
index.add(text_embeddings)

Paso 4: Retrieval – búsqueda de fragmentos relevantes

Para identificar en la base de datos vectorial aquellos fragmentos más similares a la pregunta que enviaremos al modelo, realizamos una búsqueda con la función index.search, que toma como argumentos la pregunta (previamente vectorizada) y el número de vectores similares a recuperar k. Esta función devuelve las distancias y los índices de los vectores más similares, y con estos índices se recuperan los fragmentos de texto más próximos en significado a nuestra pregunta.

Hay varias estrategias de recuperación: además de la búsqueda por similitud, se pueden usar metadatos o métodos estadísticos como TF-IDF y BM25. A veces, es útil recuperar fragmentos más grandes que incluyan contexto adicional. Un problema común en estos casos es que la información en el medio de un contexto largo se pierde, el conocido como “lost in the middle problem”. Por eso, en contextos muy largos se recomienda experimentar con el orden de los fragmentos recuperados, para ver cual nos puede dar un prompt más óptimo.

Finalmente enviaremos estos fragmentos al modelo junto a la pregunta que le queremos hacer para condicionar así al modelo.

Paso 5: Recuperación y generación de respuestas

Abajo mostramos el código cuando enviamos al modelo la pregunta “”What was Achilles’ ending?”:

# Creamos los embeddings para la pregunta
question = "What was Achilles’ ending?"
question_embeddings = np.array([get_text_embedding(question)])

# Buscamos los dos (k=2) fragmentos más similares en la base de datos vectorial
D, I = index.search(question_embeddings, k=2)
retrieved_chunk = [chunks[i] for i in I.tolist()[0]]

# Generamos una respuesta utilizando los fragmentos recuperados como contexto
prompt1 = f"""
Context information is below.
---------------------
{retrieved_chunk}
---------------------
Given the context information and not prior knowledge, answer the query.
Query: {question}
Answer:
"""

def run_mistral(user_message, model="mistral-medium-latest"):
    messages = [
        {"role": "user", "content": user_message}
    ]
    chat_response = client.chat(
        model=model,
        messages=messages
    )
    return chat_response.choices[0].message.content

A continuación comparamos las repuestas generadas usando RAG vs. las respuesta generadas por el LLM de Mistral sin RAG:

#Repuesta con RAG:
run_mistral(prompt1)
output: 
"According to the provided context, Achilles was given a second chance at life by Zeus. After leaving the underworld, he spent the rest of his life as a shepherd in the hills of Ithaca."

#Repuesta plain model (sin RAG):
run_mistral(prompt2)
output: 
"In Greek mythology, Achilles was killed by an arrow shot by Paris that struck his one vulnerable spot, his heel."

Resultados con RAG y conclusiones

Si hiciéramos nuestra prueba de RAG con la versión original de la Odisea veríamos que las respuestas generadas con y sin RAG mostrarían poca diferencia porque el modelo de Mistral ya ha sido entrenado con el texto de la Odisea. La diferencia principal sería que las respuestas generadas con RAG tienden a ser más directas y cortas, ya que se basan únicamente en el contexto y el prompt proporcionados.

En nuestro ejemplo, vemos como RAG mejora la capacidad de un modelo de lenguaje para incorporar y responder basándose en nueva información. Hemos utilizado una versión modificada de “La Odisea”, cambiando el desenlace de la historia de Ulises, para evaluar la adaptabilidad de los modelos. Observamos que cuando usamos RAG y le presentamos el texto modificado, el modelo genera respuestas que reflejan este nuevo final, demostrando que RAG efectivamente condiciona al modelo para utilizar información actualizada. En contraste, al consultar al modelo estándar, sin RAG, éste solo puede responder con el conocimiento previo con el que fue entrenado, reiterando el final original.

Implementar RAG en Mistral es una solución potente y flexible para mejorar la precisión y relevancia de las respuestas generadas por modelos de lenguaje. Mistral facilita esta implementación de forma open-source, permitiendo personalizar y ajustar tu sistema manteniendo el control de los datos y la privacidad de estos.

*RAG vs fine tuning

En nuestro ejemplo utilizamos RAG para ajustar las respuestas de un modelo de lenguaje con información reciente. Pero otra técnica que podría haberse empleado es el Fine Tuning. A diferencia del RAG, que integra información nueva sin modificar el modelo, el Fine Tuning implica reentrenar el modelo con nuevos datos, modificando sus parámetros internos.

El Fine Tuning es más efectivo en contextos donde la precisión y profundidad del conocimiento son críticos, permitiendo que el modelo refleje de forma más completa y permanente la información actualizada.

Por otro lado RAG es una alternativa eficiente y menos costosa que permite integrar información actualizada en tiempo real. Al combinar la generación de texto con la recuperación de datos relevantes, RAG se puede adaptar a diversas aplicaciones que pueden requerir información actualizada.

Keep reading