Construir RAG na arquitetura Arm

As CPUsArm são amplamente utilizadas em uma ampla gama de aplicações, incluindo casos de uso tradicionais de aprendizado de máquina (ML) e inteligência artificial (AI).

Neste tutorial, aprende a construir uma aplicação Retrieval-Augmented Generation (RAG) em infra-estruturas baseadas em Arm. Para o armazenamento de vetores, utilizamos o Zilliz Cloud, o banco de dados de vetores Milvus totalmente gerenciado. O Zilliz Cloud está disponível nas principais nuvens, como AWS, GCP e Azure. Nesta demonstração, usamos o Zilliz Cloud implantado na AWS com máquinas Arm. Para LLM, usamos o modelo Llama-3.1-8B na CPU do servidor baseado em Arm da AWS usando llama.cpp.

Pré-requisitos

Para executar este exemplo, recomendamos a utilização do AWS Graviton, que fornece uma forma económica de executar cargas de trabalho de ML em servidores baseados em Arm. Este notebook foi testado em uma instância do AWS Graviton3 c7g.2xlarge com o sistema Ubuntu 22.04 LTS.

Você precisa de pelo menos quatro núcleos e 8 GB de RAM para executar este exemplo. Configure o armazenamento em disco até pelo menos 32 GB. Recomendamos que você use uma instância com a mesma especificação ou melhor.

Depois de iniciar a instância, ligue-se a ela e execute os seguintes comandos para preparar o ambiente.

Instalar o python no servidor:

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

Criar e ativar um ambiente virtual:

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

Instalar as dependências python necessárias:

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

Carregamento de dados offline

Criar a coleção

Utilizamos o Zilliz Cloud implementado no AWS com máquinas baseadas em Arm para armazenar e recuperar os dados vectoriais. Para começar rapidamente, basta registar uma conta no Zilliz Cloud gratuitamente.

Definimos uri e token como o ponto de extremidade público e a chave Api no Zilliz Cloud.

from pymilvus import MilvusClient

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

collection_name = "my_rag_collection"

Verifique se a coleção já existe e elimine-a se existir.

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

Crie uma nova coleção com os parâmetros especificados.

Se não especificarmos qualquer informação de campo, o Milvus criará automaticamente um campo id por defeito para a chave primária e um campo vector para armazenar os dados do vetor. Um campo JSON reservado é utilizado para armazenar campos não definidos pelo esquema e os respectivos valores.

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

Utilizamos a distância do produto interno como o tipo de métrica predefinido. Para obter mais informações sobre os tipos de distância, pode consultar a página Métricas de semelhança

Preparar os dados

Utilizamos as páginas de FAQ da Documentação do Milvus 2.4.x como conhecimento privado no nosso RAG, que é uma boa fonte de dados para um pipeline RAG simples.

Descarregue o ficheiro zip e extraia os documentos para a pasta 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

Carregamos todos os ficheiros markdown da pasta milvus_docs/en/faq. Para cada documento, utilizamos simplesmente "#" para separar o conteúdo do ficheiro, o que permite separar aproximadamente o conteúdo de cada parte principal do ficheiro 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("# ")

Inserir dados

Preparamos um modelo de incorporação simples mas eficiente, o all-MiniLM-L6-v2, que pode converter texto em vectores de incorporação.

from langchain_huggingface import HuggingFaceEmbeddings

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

Iteramos pelas linhas de texto, criamos embeddings e depois inserimos os dados no Milvus.

Aqui está um novo campo text, que é um campo não definido no esquema da coleção. Será automaticamente adicionado ao campo dinâmico JSON reservado, que pode ser tratado como um campo normal a um nível elevado.

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]

Iniciar o serviço LLM no braço

Nesta secção, vamos construir e lançar o serviço llama.cpp no CPU baseado em Arm.

Modelo Llama 3.1 & llama.cpp

O modelo Llama-3.1-8B da Meta pertence à família de modelos Llama 3.1 e é de uso livre para fins comerciais e de pesquisa. Antes de utilizar o modelo, visite o sítio Web da Llama e preencha o formulário para pedir acesso.

O llama.cpp é um projeto C/C++ de código aberto que permite uma inferência LLM eficiente numa variedade de hardware - tanto localmente como na nuvem. É possível hospedar convenientemente um modelo Llama 3.1 usando llama.cpp.

Baixe e construa llama.cpp

Execute os seguintes comandos para instalar make, cmake, gcc, g++ e outras ferramentas essenciais necessárias para compilar llama.cpp a partir da fonte:

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

Agora você está pronto para começar a compilar llama.cpp.

Clone o repositório de código-fonte para llama.cpp:

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

Por padrão, llama.cpp constrói para CPU apenas no Linux e no Windows. Não precisa de fornecer quaisquer interruptores extra para o construir para o CPU Arm em que o executa.

Execute make para compilá-lo:

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

Verifique se llama.cpp foi compilado corretamente executando o comando help:

$ ./llama-cli -h

Se llama.cpp foi compilado corretamente, você verá a opção help exibida. O snippet de saída tem o seguinte aspeto:

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

Pode agora descarregar o modelo utilizando o 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

O formato do modelo GGUF, introduzido pela equipa llama.cpp, usa compressão e quantização para reduzir a precisão dos pesos para inteiros de 4 bits, diminuindo significativamente as exigências computacionais e de memória e tornando as CPUs Arm eficazes para a inferência LLM.

Re-quantificar os pesos do modelo

Para re-quantizar, execute

$ ./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

Isto produzirá um novo ficheiro, dolphin-2.9.4-llama3.1-8b-Q4_0_8_8.gguf, que contém pesos reconfigurados que permitem a llama-cli utilizar o suporte SVE 256 e MATMUL_INT8.

Iniciar o serviço LLM

Pode utilizar o programa servidor llama.cpp e enviar pedidos através de uma API compatível com OpenAI. Isto permite-lhe desenvolver aplicações que interagem com o LLM várias vezes sem ter de o iniciar e parar repetidamente. Adicionalmente, pode aceder ao servidor a partir de outra máquina onde o LLM está alojado através da rede.

Inicie o servidor a partir da linha de comando, e ele escuta na porta 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

Também é possível ajustar os parâmetros do LLM iniciado para adaptá-lo ao hardware do servidor e obter o desempenho ideal. Para obter mais informações sobre os parâmetros, consulte o comando llama-server --help.

Se tiver dificuldade em executar esta etapa, pode consultar os documentos oficiais para obter mais informações.

Iniciou o serviço LLM na sua CPU baseada em Arm. Em seguida, interagimos diretamente com o serviço usando o OpenAI SDK.

RAG online

Cliente LLM e modelo de incorporação

Inicializamos o cliente LLM e preparamos o modelo de incorporação.

Para o LLM, usamos o OpenAI SDK para solicitar o serviço Llama lançado anteriormente. Não precisamos de usar nenhuma chave de API porque, na verdade, é o nosso serviço local llama.cpp.

from openai import OpenAI

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

Gerar um embedding de teste e imprimir a sua dimensão e os primeiros 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 dados para uma consulta

Vamos especificar uma pergunta frequente sobre o Milvus.

question = "How is data stored in milvus?"

Pesquise a pergunta na coleção e obtenha as 3 melhores correspondências 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
)

Vamos dar uma olhadela aos resultados da pesquisa da 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 o LLM para obter uma resposta RAG

Converta os documentos recuperados num formato de cadeia de caracteres.

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>
"""

Use o LLM para gerar uma resposta com base nos prompts. Definimos o parâmetro model como not-used, pois é um parâmetro redundante para o serviço 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.

Parabéns! Construiu uma aplicação RAG sobre as infra-estruturas baseadas em Arm.