Costruire RAG con Milvus e Fireworks AI

Fireworks AI è una piattaforma di inferenza AI generativa che offre una velocità e una prontezza di produzione leader nel settore per l'esecuzione e la personalizzazione dei modelli. Fireworks AI fornisce una serie di servizi di AI generativa, tra cui modelli serverless, distribuzioni on-demand e capacità di messa a punto. Offre un ambiente completo per la distribuzione di vari modelli di IA, compresi i modelli linguistici di grandi dimensioni (LLM) e i modelli di incorporazione. Fireworks AI aggrega numerosi modelli, consentendo agli utenti di accedere e utilizzare facilmente queste risorse senza la necessità di configurare un'infrastruttura estesa.

In questo tutorial vi mostreremo come costruire una pipeline RAG (Retrieval-Augmented Generation) con Milvus e Fireworks AI.

Preparazione

Dipendenze e ambiente

$ pip install --upgrade pymilvus openai requests tqdm

Fireworks AI abilita l'API in stile OpenAI. È possibile accedere al suo sito web ufficiale e preparare la chiave api FIREWORKS_API_KEY come variabile d'ambiente.

import os

os.environ["FIREWORKS_API_KEY"] = "***********"

Preparare i dati

Come conoscenza privata nel nostro RAG utilizziamo le pagine FAQ della Documentazione Milvus 2.4.x, che è una buona fonte di dati per una semplice pipeline RAG.

Scaricare il file zip ed estrarre i documenti nella cartella 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

Carichiamo tutti i file markdown dalla cartella milvus_docs/en/faq. Per ogni documento, usiamo semplicemente "# " per separare il contenuto del file, che può separare approssimativamente il contenuto di ogni parte principale del file 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("# ")

Preparare l'LLM e il modello di incorporamento

Inizializziamo un client per preparare l'LLM e il modello di incorporamento. Fireworks AI abilita l'API in stile OpenAI e si può usare la stessa API con piccoli aggiustamenti per chiamare il modello di embedding e l'LLM.

from openai import OpenAI

fireworks_client = OpenAI(
    api_key=os.environ["FIREWORKS_API_KEY"],
    base_url="https://api.fireworks.ai/inference/v1",
)

Definire una funzione per generare embeddings di testo utilizzando il client. Utilizziamo il modello nomic-ai/nomic-embed-text-v1.5 come esempio.

def emb_text(text):
    return (
        fireworks_client.embeddings.create(
            input=text, model="nomic-ai/nomic-embed-text-v1.5"
        )
        .data[0]
        .embedding
    )

Generare un embedding di prova e stamparne la dimensione e i primi elementi.

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

768
[0.04815673828125, 0.0261993408203125, -0.1749267578125, -0.03131103515625, 0.068115234375, -0.00621795654296875, 0.03955078125, -0.0210723876953125, 0.039703369140625, -0.0286102294921875]

Caricare i dati in Milvus

Creare la raccolta

from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="./milvus_demo.db")

collection_name = "my_rag_collection"

Verificare se la raccolta esiste già e, in caso affermativo, eliminarla.

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

Creare una nuova raccolta con i parametri specificati.

Se non si specifica alcun campo, Milvus creerà automaticamente un campo predefinito id per la chiave primaria e un campo vector per memorizzare i dati vettoriali. Un campo JSON riservato viene utilizzato per memorizzare campi non definiti dalla mappa e i loro valori.

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

Inserire i dati

Si intersecano le righe di testo, si creano le incorporazioni e si inseriscono i dati in Milvus.

Ecco un nuovo campo text, che è un campo non definito nello schema della collezione. Verrà aggiunto automaticamente al campo dinamico JSON riservato, che può essere trattato come un campo normale ad alto livello.

from tqdm import tqdm

data = []

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

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

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





{'insert_count': 72, 'ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71], 'cost': 0}

Costruire la RAG

Recuperare i dati per una query

Specifichiamo una domanda frequente su Milvus.

question = "How is data stored in milvus?"

Cerchiamo la domanda nella raccolta e recuperiamo le prime tre corrispondenze semantiche.

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[
        emb_text(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
)

Diamo un'occhiata ai risultati della ricerca della query

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.8334928750991821
    ],
    [
        "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.746377170085907
    ],
    [
        "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.7328270673751831
    ]
]

Utilizzare LLM per ottenere una risposta RAG

Convertire i documenti recuperati in un formato stringa.

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

Definire i prompt del sistema e dell'utente per il Lanage Model. Questo prompt viene assemblato con i documenti recuperati da 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>
"""

Utilizzare il modello llama-v3p1-405b-instruct fornito da Fireworks per generare una risposta basata sui prompt.

response = fireworks_client.chat.completions.create(
    model="accounts/fireworks/models/llama-v3p1-405b-instruct",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT},
    ],
)
print(response.choices[0].message.content)

According to the provided context, Milvus stores data in two ways:

1. Inserted data, including vector data, scalar data, and collection-specific schema, are stored in persistent storage as incremental logs. This can be done using multiple object storage backends such as MinIO, AWS S3, Google Cloud Storage, Azure Blob Storage, Alibaba Cloud OSS, and Tencent Cloud Object Storage.
2. Metadata, which are generated within Milvus, are stored in etcd, with each Milvus module having its own metadata.

Additionally, when data is inserted, it is first loaded into a message queue, and then written to persistent storage as incremental logs by the data node. The `flush()` function can be used to force the data node to write all data in the message queue to persistent storage immediately.

Ottimo! Abbiamo costruito con successo una pipeline RAG con Milvus e Fireworks AI.