Costruire RAG con Milvus e Gemini

L'API Gemini e Google AI Studio consentono di iniziare a lavorare con i modelli più recenti di Google e di trasformare le proprie idee in applicazioni scalabili. Gemini fornisce l'accesso a potenti modelli linguistici come Gemini-1.5-Flash, Gemini-1.5-Flash-8B e Gemini-1.5-Pro per attività come la generazione di testi, l'elaborazione di documenti, la visione, l'analisi audio e altro ancora. L'API consente di inserire contesti lunghi con milioni di token, di mettere a punto i modelli per compiti specifici, di generare output strutturati come JSON e di sfruttare funzionalità come il recupero semantico e l'esecuzione di codice.

In questo tutorial vi mostreremo come costruire una pipeline RAG (Retrieval-Augmented Generation) con Milvus e Gemini. Utilizzeremo il modello Gemini per generare testo in base a una determinata query. Utilizzeremo anche Milvus per memorizzare e recuperare il testo generato.

Preparazione

Dipendenze e ambiente

$ pip install --upgrade pymilvus google-generativeai requests tqdm

Per prima cosa è necessario accedere alla piattaforma Google AI Studio e preparare la chiave api GEMINI_API_KEY come variabile d'ambiente.

import os

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

Preparare i dati

Nel nostro RAG utilizziamo le pagine FAQ della Documentazione Milvus 2.4.x come conoscenza privata, 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

Utilizziamo gemini-1.5-flash come LLM e text-embedding-004 come modello di incorporamento.

Proviamo a generare una risposta di prova dall'LLM:

import google.generativeai as genai

genai.configure(api_key=os.environ["GEMINI_API_KEY"])

gemini_model = genai.GenerativeModel("gemini-1.5-flash")

response = gemini_model.generate_content("who are you")
print(response.text)

I am a large language model, trained by Google.  I am an AI and don't have a personal identity or consciousness.  My purpose is to process information and respond to a wide range of prompts and questions in a helpful and informative way.

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

test_embeddings = genai.embed_content(
    model="models/text-embedding-004", content=["This is a test1", "This is a test2"]
)["embedding"]

embedding_dim = len(test_embeddings[0])
print(embedding_dim)
print(test_embeddings[0][:10])

768
[0.013588584, -0.004361838, -0.08481652, -0.039724775, 0.04723794, -0.0051557426, 0.026071774, 0.045514572, -0.016867816, 0.039378334]

Caricare i dati in Milvus

Creare la collezione

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 da schemi 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 = []

doc_embeddings = genai.embed_content(
    model="models/text-embedding-004", content=text_lines
)["embedding"]

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

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

Creating embeddings: 100%|██████████| 72/72 [00:00<00:00, 468201.38it/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 3 corrispondenze semantiche.

question_embedding = genai.embed_content(
    model="models/text-embedding-004", content=question
)["embedding"]

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[question_embedding],
    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 domanda

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.8048275113105774
    ],
    [
        "Does the query perform in memory? What are incremental data and historical data?\n\nYes. When a query request comes, Milvus searches both incremental data and historical data by loading them into memory. Incremental data are in the growing segments, which are buffered in memory before they reach the threshold to be persisted in storage engine, while historical data are from the sealed segments that are stored in the object storage. Incremental data and historical data together constitute the whole dataset to search.\n\n###",
        0.7574886679649353
    ],
    [
        "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.7453608512878418
    ]
]

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 Gemini per generare una risposta basata sui prompt.

gemini_model = genai.GenerativeModel(
    "gemini-1.5-flash", system_instruction=SYSTEM_PROMPT
)
response = gemini_model.generate_content(USER_PROMPT)
print(response.text)

Milvus stores data in two ways:  Inserted data (vector data, scalar data, and collection-specific schema) is stored as an incremental log in persistent storage using object storage backends such as MinIO, AWS S3, Google Cloud Storage, Azure Blob Storage, Alibaba Cloud OSS, and Tencent Cloud Object Storage.  Metadata, generated by each Milvus module, is stored in etcd.

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