Construir RAG en arquitectura Arm

Las CPUsArm se utilizan ampliamente en una amplia gama de aplicaciones, incluyendo casos de uso tradicionales de aprendizaje automático (ML) e inteligencia artificial (AI).

En este tutorial, aprenderá a crear una aplicación de generación mejorada de recuperación (RAG) en infraestructuras basadas en Arm. Para el almacenamiento vectorial, utilizamos Zilliz Cloud, la base de datos vectorial Milvus totalmente gestionada. Zilliz Cloud está disponible en las principales nubes como AWS, GCP y Azure. En esta demo utilizamos Zilliz Cloud desplegado en AWS con máquinas Arm. Para LLM, utilizamos el modelo Llama-3.1-8B en la CPU del servidor basado en Arm de AWS utilizando llama.cpp.

Requisitos previos

Para ejecutar este ejemplo, recomendamos utilizar AWS Graviton, que proporciona una forma rentable de ejecutar cargas de trabajo ML en servidores basados en Arm. Este cuaderno se ha probado en una instancia de AWS Graviton3 c7g.2xlarge con el sistema Ubuntu 22.04 LTS.

Necesita al menos cuatro núcleos y 8 GB de RAM para ejecutar este ejemplo. Configure un almacenamiento en disco de al menos 32 GB. Le recomendamos que utilice una instancia con las mismas o mejores especificaciones.

Después de lanzar la instancia, conéctese a ella y ejecute los siguientes comandos para preparar el entorno.

Instale python en el servidor:

$ sudo apt update
$ sudo apt install python-is-python3 python3-pip python3-venv -y

Cree y active un entorno virtual:

$ python -m venv venv
$ source venv/bin/activate

Instalar las dependencias de python necesarias:

$ pip install --upgrade pymilvus openai requests langchain-huggingface huggingface_hub tqdm

Carga de datos sin conexión

Crear la colección

Utilizamos Zilliz Cloud desplegado en AWS con máquinas basadas en Arm para almacenar y recuperar los datos vectoriales. Para empezar rápidamente, basta con registrar una cuenta en Zilliz Cloud de forma gratuita.

Configuramos uri y token como Public Endpoint y Api key en Zilliz Cloud.

from pymilvus import MilvusClient

milvus_client = MilvusClient(
    uri="<your_zilliz_public_endpoint>", token="<your_zilliz_api_key>"
)

collection_name = "my_rag_collection"

Compruebe si la colección ya existe y elimínela en caso afirmativo.

if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

Creamos una nueva colección con los parámetros especificados.

Si no especificamos ninguna información de campo, Milvus creará automáticamente un campo id por defecto para la clave primaria, y un campo vector para almacenar los datos vectoriales. Se utiliza un campo JSON reservado para almacenar campos no definidos por el esquema y sus valores.

milvus_client.create_collection(
    collection_name=collection_name,
    dimension=384,
    metric_type="IP",  # Inner product distance
    consistency_level="Strong",  # Strong consistency level
)

Utilizamos la distancia producto interior como tipo de métrica por defecto. Para más información sobre los tipos de distancia, puede consultar la página Métricas de similitud.

Preparar los datos

Utilizamos las páginas FAQ de la Documentación de Milvus 2.4.x como conocimiento privado en nuestra RAG, que es una buena fuente de datos para una canalización RAG sencilla.

Descargamos el archivo zip y extraemos los documentos a la carpeta milvus_docs.

$ wget https://github.com/milvus-io/milvus-docs/releases/download/v2.4.6-preview/milvus_docs_2.4.x_en.zip
$ unzip -q milvus_docs_2.4.x_en.zip -d milvus_docs

Cargamos todos los archivos markdown de la carpeta milvus_docs/en/faq. Para cada documento, simplemente usamos "# " para separar el contenido en el archivo, lo que puede separar aproximadamente el contenido de cada parte principal del archivo markdown.

from glob import glob

text_lines = []

for file_path in glob("milvus_docs/en/faq/*.md", recursive=True):
    with open(file_path, "r") as file:
        file_text = file.read()

    text_lines += file_text.split("# ")

Insertar datos

Preparamos un modelo de incrustación simple pero eficiente all-MiniLM-L6-v2 que puede convertir texto en vectores de incrustación.

from langchain_huggingface import HuggingFaceEmbeddings

embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

Iteramos a través de las líneas de texto, creamos incrustaciones y luego insertamos los datos en Milvus.

Aquí hay un nuevo campo text, que es un campo no definido en el esquema de la colección. Se añadirá automáticamente al campo dinámico JSON reservado, que puede tratarse como un campo normal a alto nivel.

from tqdm import tqdm

data = []

text_embeddings = embedding_model.embed_documents(text_lines)

for i, (line, embedding) in enumerate(
    tqdm(zip(text_lines, text_embeddings), desc="Creating embeddings")
):
    data.append({"id": i, "vector": embedding, "text": line})

milvus_client.insert(collection_name=collection_name, data=data)

Creating embeddings: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 72/72 [00:18<00:00,  3.91it/s]

Lanzar el servicio LLM en Arm

En esta sección, construiremos y lanzaremos el servicio llama.cpp en la CPU basada en Arm.

Modelo Llama 3.1 y llama.cpp

El modelo Llama-3.1-8B de Meta pertenece a la familia de modelos Llama 3.1 y es de uso libre para investigación y fines comerciales. Antes de utilizar el modelo, visite el sitio web de Llama y rellene el formulario para solicitar acceso.

llama.cpp es un proyecto C/C++ de código abierto que permite una inferencia LLM eficiente en una variedad de hardware, tanto localmente como en la nube. Puede alojar cómodamente un modelo Llama 3.1 utilizando llama.cpp.

Descarga y compila llama.cpp

Ejecuta los siguientes comandos para instalar make, cmake, gcc, g++ y otras herramientas esenciales necesarias para compilar llama.cpp desde el código fuente:

$ sudo apt install make cmake -y
$ sudo apt install gcc g++ -y
$ sudo apt install build-essential -y

Ahora estás listo para empezar a construir llama.cpp.

Clona el repositorio fuente de llama.cpp:

$ git clone https://github.com/ggerganov/llama.cpp

Por defecto, llama.cpp construye sólo para CPU en Linux y Windows. No necesitas proporcionar ningún interruptor extra para construirlo para la CPU Arm en la que lo ejecutes.

Ejecute make para compilarlo:

$ cd llama.cpp
$ make GGML_NO_LLAMAFILE=1 -j$(nproc)

Comprueba que llama.cpp se ha compilado correctamente ejecutando el comando help:

$ ./llama-cli -h

Si llama.cpp se ha compilado correctamente, verás que aparece la opción help. El fragmento de salida tiene este aspecto:

example usage:

    text generation:     ./llama-cli -m your_model.gguf -p "I believe the meaning of life is" -n 128

    chat (conversation): ./llama-cli -m your_model.gguf -p "You are a helpful assistant" -cnv

Ahora puede descargar el modelo utilizando el cli huggingface:

$ huggingface-cli download cognitivecomputations/dolphin-2.9.4-llama3.1-8b-gguf dolphin-2.9.4-llama3.1-8b-Q4_0.gguf --local-dir . --local-dir-use-symlinks False

El formato del modelo GGUF, introducido por el equipo de llama.cpp, utiliza compresión y cuantización para reducir la precisión de los pesos a enteros de 4 bits, disminuyendo significativamente las demandas computacionales y de memoria y haciendo que las CPUs Arm sean efectivas para la inferencia LLM.

Recuantificar los pesos del modelo

Para volver a cuantificar, ejecute

$ ./llama-quantize --allow-requantize dolphin-2.9.4-llama3.1-8b-Q4_0.gguf dolphin-2.9.4-llama3.1-8b-Q4_0_8_8.gguf Q4_0_8_8

Esto generará un nuevo archivo, dolphin-2.9.4-llama3.1-8b-Q4_0_8_8.gguf, que contiene pesos reconfigurados que permiten a llama-cli utilizar SVE 256 y el soporte MATMUL_INT8.

Iniciar el servicio LLM

Puedes utilizar el programa servidor llama.cpp y enviar peticiones a través de una API compatible con OpenAI. Esto te permite desarrollar aplicaciones que interactúen con el LLM varias veces sin tener que iniciarlo y detenerlo repetidamente. Además, puedes acceder al servidor desde otra máquina en la que el LLM esté alojado a través de la red.

Inicia el servidor desde la línea de comandos, y escucha en el puerto 8080:

$ ./llama-server -m dolphin-2.9.4-llama3.1-8b-Q4_0_8_8.gguf -n 2048 -t 64 -c 65536  --port 8080

'main: server is listening on 127.0.0.1:8080 - starting the main loop

También puede ajustar los parámetros del LLM iniciado para adaptarlo al hardware de su servidor y obtener un rendimiento ideal. Para obtener más información sobre los parámetros, consulta el comando llama-server --help.

Si tienes dificultades para realizar este paso, puedes consultar los documentos oficiales para obtener más información.

Ya has iniciado el servicio LLM en tu CPU basada en Arm. A continuación, interactuaremos directamente con el servicio utilizando el SDK de OpenAI.

RAG en línea

Cliente LLM y modelo de incrustación

Inicializamos el cliente LLM y preparamos el modelo de incrustación.

Para el LLM, utilizamos el SDK de OpenAI para solicitar el servicio Llama lanzado anteriormente. No necesitamos usar ninguna clave API porque en realidad es nuestro servicio local llama.cpp.

from openai import OpenAI

llm_client = OpenAI(base_url="http://localhost:8080/v1", api_key="no-key")

Generamos una incrustación de prueba e imprimimos su dimensión y sus primeros elementos.

test_embedding = embedding_model.embed_query("This is a test")
embedding_dim = len(test_embedding)
print(embedding_dim)
print(test_embedding[:10])

384
[0.03061249852180481, 0.013831384479999542, -0.02084377221763134, 0.016327863559126854, -0.010231520049273968, -0.0479842908680439, -0.017313342541456223, 0.03728749603033066, 0.04588735103607178, 0.034405000507831573]

Recuperar datos para una consulta

Especifiquemos una pregunta frecuente sobre Milvus.

question = "How is data stored in milvus?"

Busquemos la pregunta en la colección y recuperemos las 3 primeras coincidencias semánticas.

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[
        embedding_model.embed_query(question)
    ],  # Use the `emb_text` function to convert the question to an embedding vector
    limit=3,  # Return top 3 results
    search_params={"metric_type": "IP", "params": {}},  # Inner product distance
    output_fields=["text"],  # Return the text field
)

Veamos los resultados de la consulta

import json

retrieved_lines_with_distances = [
    (res["entity"]["text"], res["distance"]) for res in search_res[0]
]
print(json.dumps(retrieved_lines_with_distances, indent=4))

[
    [
        " Where does Milvus store data?\n\nMilvus deals with two types of data, inserted data and metadata. \n\nInserted data, including vector data, scalar data, and collection-specific schema, are stored in persistent storage as incremental log. Milvus supports multiple object storage backends, including [MinIO](https://min.io/), [AWS S3](https://aws.amazon.com/s3/?nc1=h_ls), [Google Cloud Storage](https://cloud.google.com/storage?hl=en#object-storage-for-companies-of-all-sizes) (GCS), [Azure Blob Storage](https://azure.microsoft.com/en-us/products/storage/blobs), [Alibaba Cloud OSS](https://www.alibabacloud.com/product/object-storage-service), and [Tencent Cloud Object Storage](https://www.tencentcloud.com/products/cos) (COS).\n\nMetadata are generated within Milvus. Each Milvus module has its own metadata that are stored in etcd.\n\n###",
        0.6488019824028015
    ],
    [
        "How does Milvus flush data?\n\nMilvus returns success when inserted data are loaded to the message queue. However, the data are not yet flushed to the disk. Then Milvus' data node writes the data in the message queue to persistent storage as incremental logs. If `flush()` is called, the data node is forced to write all data in the message queue to persistent storage immediately.\n\n###",
        0.5974207520484924
    ],
    [
        "What is the maximum dataset size Milvus can handle?\n\n  \nTheoretically, the maximum dataset size Milvus can handle is determined by the hardware it is run on, specifically system memory and storage:\n\n- Milvus loads all specified collections and partitions into memory before running queries. Therefore, memory size determines the maximum amount of data Milvus can query.\n- When new entities and and collection-related schema (currently only MinIO is supported for data persistence) are added to Milvus, system storage determines the maximum allowable size of inserted data.\n\n###",
        0.5833579301834106
    ]
]

Utilizar LLM para obtener una respuesta RAG

Convertir los documentos recuperados en un formato de cadena.

context = "\n".join(
    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]
)

Define system and user prompts for the Language Model. This prompt is assembled with the retrieved documents from Milvus.

SYSTEM_PROMPT = """
Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided.
"""
USER_PROMPT = f"""
Use the following pieces of information enclosed in <context> tags to provide an answer to the question enclosed in <question> tags.
<context>
{context}
</context>
<question>
{question}
</question>
"""

Utiliza LLM para generar una respuesta basada en las instrucciones. Establecemos el parámetro model en not-used ya que es un parámetro redundante para el servicio llama.cpp.

response = llm_client.chat.completions.create(
    model="not-used",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT},
    ],
)
print(response.choices[0].message.content)

Milvus stores data in two types: inserted data and metadata. Inserted data, including vector data, scalar data, and collection-specific schema, are stored in persistent storage as incremental log. Milvus supports multiple object storage backends such as MinIO, AWS S3, Google Cloud Storage (GCS), Azure Blob Storage, Alibaba Cloud OSS, and Tencent Cloud Object Storage (COS). Metadata are generated within Milvus and each Milvus module has its own metadata that are stored in etcd.

¡Enhorabuena! Has construido una aplicación RAG sobre las infraestructuras basadas en Arm.