RAG mit Milvus und Gemini aufbauen

Die Gemini-API und Google AI Studio helfen Ihnen, mit den neuesten Modellen von Google zu arbeiten und Ihre Ideen in skalierbare Anwendungen umzusetzen. Gemini bietet Zugang zu leistungsstarken Sprachmodellen wie Gemini-1.5-Flash, Gemini-1.5-Flash-8B und Gemini-1.5-Pro für Aufgaben wie Texterstellung, Dokumentenverarbeitung, Bildverarbeitung, Audioanalyse und mehr. Die API ermöglicht die Eingabe langer Kontexte mit Millionen von Token, die Feinabstimmung von Modellen für bestimmte Aufgaben, die Erzeugung strukturierter Ausgaben wie JSON und die Nutzung von Funktionen wie semantisches Retrieval und Codeausführung.

In diesem Tutorial zeigen wir Ihnen, wie Sie eine RAG-Pipeline (Retrieval-Augmented Generation) mit Milvus und Gemini aufbauen. Wir werden das Gemini-Modell verwenden, um Text auf der Grundlage einer gegebenen Anfrage zu generieren. Wir werden auch Milvus verwenden, um den generierten Text zu speichern und abzurufen.

Vorbereitung

Abhängigkeiten und Umgebung

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

Sie sollten sich zunächst bei der Google AI Studio-Plattform anmelden und den api-Schlüssel GEMINI_API_KEY als Umgebungsvariable vorbereiten.

import os

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

Bereiten Sie die Daten vor

Wir verwenden die FAQ-Seiten aus der Milvus-Dokumentation 2.4.x als privates Wissen in unserem RAG, was eine gute Datenquelle für eine einfache RAG-Pipeline ist.

Laden Sie die Zip-Datei herunter und entpacken Sie die Dokumente in den Ordner 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

Wir laden alle Markdown-Dateien aus dem Ordner milvus_docs/en/faq. Für jedes Dokument verwenden wir einfach "# ", um den Inhalt in der Datei zu trennen, wodurch der Inhalt jedes Hauptteils der Markdown-Datei grob getrennt werden kann.

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("# ")

Vorbereiten des LLM und des Einbettungsmodells

Wir verwenden die gemini-1.5-flash als LLM und die text-embedding-004 als Einbettungsmodell.

Lassen Sie uns versuchen, eine Testantwort aus der LLM zu generieren:

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.

Erzeugen Sie eine Testeinbettung und drucken Sie ihre Dimension und die ersten Elemente aus.

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]

Laden Sie Daten in Milvus

Erstellen Sie die Sammlung

from pymilvus import MilvusClient

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

collection_name = "my_rag_collection"

Prüfen Sie, ob die Sammlung bereits existiert und löschen Sie sie, wenn dies der Fall ist.

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

Erstellen Sie eine neue Sammlung mit den angegebenen Parametern.

Wenn wir keine Feldinformationen angeben, erstellt Milvus automatisch ein Standardfeld id für den Primärschlüssel und ein Feld vector zum Speichern der Vektordaten. Ein reserviertes JSON-Feld wird verwendet, um nicht schema-definierte Felder und ihre Werte zu speichern.

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

Daten einfügen

Iterieren Sie durch die Textzeilen, erstellen Sie Einbettungen und fügen Sie dann die Daten in Milvus ein.

Hier ist ein neues Feld text, das ein nicht definiertes Feld im Sammelschema ist. Es wird automatisch dem reservierten dynamischen JSON-Feld hinzugefügt, das auf hoher Ebene wie ein normales Feld behandelt werden kann.

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}

RAG erstellen

Abrufen von Daten für eine Abfrage

Lassen Sie uns eine häufige Frage über Milvus angeben.

question = "How is data stored in milvus?"

Suchen Sie nach der Frage in der Sammlung und rufen Sie die semantischen Top-3-Treffer ab.

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
)

Werfen wir einen Blick auf die Suchergebnisse der Abfrage

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

LLM verwenden, um eine RAG-Antwort zu erhalten

Konvertieren Sie die abgerufenen Dokumente in ein String-Format.

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

Definieren Sie System- und Benutzer-Prompts für das Lanage Model. Diese Eingabeaufforderung wird mit den abgerufenen Dokumenten aus Milvus zusammengestellt.

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

Verwenden Sie den Gemini, um eine Antwort auf der Grundlage der Prompts zu generieren.

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.

Großartig! Wir haben erfolgreich eine RAG-Pipeline mit Milvus und Gemini aufgebaut.