MilvusとLepton AIでRAGを構築する

Lepton AIは、開発者や企業がAIアプリケーションを数分で効率的に実行し、本番環境で利用可能な規模にすることを可能にします。 Lepton AIは、Pythonネイティブな方法でモデルを構築し、ローカルでモデルをデバッグおよびテストし、単一のコマンドでクラウドにデプロイし、シンプルで柔軟なAPIを使用して任意のアプリケーションでモデルを利用することができます。大規模言語モデル（LLM）や拡散モデルを含む様々なAIモデルを、大規模なインフラを構築することなくデプロイするための包括的な環境を提供します。

このチュートリアルでは、MilvusとLepton AIを使ってRAG（Retrieval-Augmented Generation）パイプラインを構築する方法を紹介します。

準備

依存関係と環境

$ pip install --upgrade pymilvus[model] openai requests tqdm

LeptonはOpenAIスタイルのAPIを有効にする。公式サイトにログインし、api key LEPTONAI_TOKEN を環境変数として用意する。

import os

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

データの準備

Milvusドキュメント2.4.xのFAQページをRAGのプライベートナレッジとして使用する。

zipファイルをダウンロードし、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

フォルダmilvus_docs/en/faq からすべてのマークダウン・ファイルをロードする。各ドキュメントについて、単に "# "を使ってファイル内のコンテンツを区切るだけで、マークダウン・ファイルの各主要部分のコンテンツを大まかに区切ることができる。

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

LLMと埋め込みモデルの準備

LeptonではOpenAIスタイルのAPIが使えますが、LLMの呼び出しにも同じAPIが使えます。

from openai import OpenAI

lepton_client = OpenAI(
    api_key=os.environ["LEPTONAI_TOKEN"],
    base_url="https://mistral-7b.lepton.run/api/v1/",
)

milvus_model を使ってテキスト埋め込みを生成するための埋め込みモデルを定義します。ここでは例として、事前に学習された軽量な埋め込みモデルであるDefaultEmbeddingFunction モデルを使います。

from pymilvus import model as milvus_model

embedding_model = milvus_model.DefaultEmbeddingFunction()

テスト埋め込みを生成し、その次元と最初の数要素を表示する。

test_embedding = embedding_model.encode_queries(["This is a test"])[0]
embedding_dim = len(test_embedding)
print(embedding_dim)
print(test_embedding[:10])

768
[-0.04836066  0.07163023 -0.01130064 -0.03789345 -0.03320649 -0.01318448
 -0.03041712 -0.02269499 -0.02317863 -0.00426028]

Milvusにデータをロードする。

コレクションの作成

from pymilvus import MilvusClient

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

collection_name = "my_rag_collection"

コレクションが既に存在するか確認し、存在する場合は削除します。

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

指定したパラメータで新しいコレクションを作成します。

フィールド情報を指定しない場合、Milvusは自動的にプライマリキー用のデフォルトid フィールドと、ベクトルデータを格納するためのvector フィールドを作成します。予約されたJSONフィールドは、スキーマで定義されていないフィールドとその値を格納するために使用されます。

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

データの挿入

テキスト行を繰り返し、エンベッディングを作成し、milvusにデータを挿入します。

ここに新しいフィールドtext 、コレクションスキーマで定義されていないフィールドです。これは予約されたJSONダイナミックフィールドに自動的に追加され、高レベルでは通常のフィールドとして扱うことができる。

from tqdm import tqdm

data = []

doc_embeddings = embedding_model.encode_documents(text_lines)

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, 1090216.20it/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の構築

クエリのデータを取得する

Milvusに関するよくある質問を指定してみましょう。

question = "How is data stored in milvus?"

コレクションで質問を検索し、セマンティックトップ3マッチを取得します。

search_res = milvus_client.search(
    collection_name=collection_name,
    data=embedding_model.encode_queries(
        [question]
    ),  # 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
)

クエリの検索結果を見てみましょう。

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.6572665572166443
    ],
    [
        "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.6312146186828613
    ],
    [
        "How does Milvus handle vector data types and precision?\n\nMilvus supports Binary, Float32, Float16, and BFloat16 vector types.\n\n- Binary vectors: Store binary data as sequences of 0s and 1s, used in image processing and information retrieval.\n- Float32 vectors: Default storage with a precision of about 7 decimal digits. Even Float64 values are stored with Float32 precision, leading to potential precision loss upon retrieval.\n- Float16 and BFloat16 vectors: Offer reduced precision and memory usage. Float16 is suitable for applications with limited bandwidth and storage, while BFloat16 balances range and efficiency, commonly used in deep learning to reduce computational requirements without significantly impacting accuracy.\n\n###",
        0.6115777492523193
    ]
]

LLMを使ってRAGレスポンスを取得する

検索されたドキュメントを文字列形式に変換する。

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

ラネージ・モデルのシステム・プロンプトとユーザー・プロンプトを定義する。このプロンプトは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>
"""

Lepton AIが提供するmistral-7b モデルを使って、プロンプトに基づいたレスポンスを生成する。

response = lepton_client.chat.completions.create(
    model="mistral-7b",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT},
    ],
)
print(response.choices[0].message.content)

Inserted data in Milvus, including vector data, scalar data, and collection-specific schema, are stored in persistent storage as incremental logs. Milvus supports multiple object storage backends such as MinIO, AWS S3, Google Cloud Storage, Azure Blob Storage, Alibaba Cloud OSS, and Tencent Cloud Object Storage (COS). Metadata are generated within Milvus and stored in etcd.

素晴らしい！MilvusとLepton AIを使ったRAGパイプラインの構築に成功した。