Construire RAG avec Milvus et Gemini

L'API Gemini et Google AI Studio vous permettent de commencer à travailler avec les derniers modèles de Google et de transformer vos idées en applications évolutives. Gemini permet d'accéder à de puissants modèles de langage tels que Gemini-1.5-Flash, Gemini-1.5-Flash-8B et Gemini-1.5-Pro pour des tâches telles que la génération de texte, le traitement de documents, la vision, l'analyse audio, etc. L'API vous permet de saisir des contextes longs avec des millions de tokens, d'affiner les modèles pour des tâches spécifiques, de générer des sorties structurées comme JSON, et d'exploiter des capacités comme la récupération sémantique et l'exécution de code.

Dans ce tutoriel, nous allons vous montrer comment construire un pipeline RAG (Retrieval-Augmented Generation) avec Milvus et Gemini. Nous utiliserons le modèle Gemini pour générer du texte sur la base d'une requête donnée. Nous utiliserons également Milvus pour stocker et récupérer le texte généré.

Préparation

Dépendances et environnement

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

Vous devez d'abord vous connecter à la plateforme Google AI Studio et préparer la clé api GEMINI_API_KEY en tant que variable d'environnement.

import os

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

Préparer les données

Nous utilisons les pages FAQ de la documentation Milvus 2.4.x comme connaissance privée dans notre RAG, ce qui constitue une bonne source de données pour un pipeline RAG simple.

Téléchargez le fichier zip et extrayez les documents dans le dossier 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

Nous chargeons tous les fichiers markdown à partir du dossier milvus_docs/en/faq. Pour chaque document, nous utilisons simplement "# " pour séparer le contenu du fichier, ce qui permet de séparer grossièrement le contenu de chaque partie principale du fichier 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("# ")

Préparation du LLM et du modèle d'intégration

Nous utilisons gemini-1.5-flash comme LLM et text-embedding-004 comme modèle d'intégration.

Essayons de générer une réponse de test à partir du 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.

Générer un embedding de test et imprimer sa dimension et ses premiers éléments.

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]

Chargement des données dans Milvus

Créer la collection

from pymilvus import MilvusClient

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

collection_name = "my_rag_collection"

Vérifier si la collection existe déjà et la supprimer si c'est le cas.

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

Créer une nouvelle collection avec les paramètres spécifiés.

Si nous ne spécifions aucune information de champ, Milvus créera automatiquement un champ id par défaut pour la clé primaire et un champ vector pour stocker les données vectorielles. Un champ JSON réservé est utilisé pour stocker les champs non définis par le schéma et leurs valeurs.

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

Insérer des données

Parcourez les lignes de texte, créez des enchâssements, puis insérez les données dans Milvus.

Voici un nouveau champ text, qui est un champ non défini dans le schéma de la collection. Il sera automatiquement ajouté au champ dynamique JSON réservé, qui peut être traité comme un champ normal à un niveau élevé.

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}

Construire un RAG

Récupérer des données pour une requête

Spécifions une question fréquente sur Milvus.

question = "How is data stored in milvus?"

Cherchons la question dans la collection et récupérons les 3 meilleures réponses sémantiques.

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
)

Jetons un coup d'œil aux résultats de la recherche de la question.

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

Utiliser LLM pour obtenir une réponse RAG

Convertir les documents récupérés dans un format de chaîne.

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

Définir les messages-guides du système et de l'utilisateur pour le modèle de langue. Cette invite est assemblée avec les documents extraits de 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>
"""

Utiliser Gemini pour générer une réponse basée sur les invites.

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.

C'est très bien ! Nous avons réussi à construire un pipeline RAG avec Milvus et Gemini.