ArmアーキテクチャでRAGを構築

ArmCPUは、従来の機械学習（ML）や人工知能（AI）のユースケースを含め、幅広いアプリケーションで幅広く利用されています。

このチュートリアルでは、Armベースのインフラストラクチャ上でRAG（Retrieval-Augmented Generation）アプリケーションを構築する方法を学びます。ベクトル・ストレージには、フルマネージドMilvusベクトル・データベースであるZilliz Cloudを利用します。Zilliz Cloudは、AWS、GCP、Azureなどの主要なクラウドで利用できる。このデモでは、ArmマシンとAWS上にデプロイされたZilliz Cloudを使用している。LLMについては、AWSのArmベースのサーバーCPU上でllama.cpp を用いてLlama-3.1-8B 。

前提条件

この例を実行するには、AWS Graviton を使用することを推奨します。AWS Graviton は、Arm ベースのサーバー上で ML ワークロードを実行するためのコスト効率の良い方法を提供します。このノートブックは、Ubuntu 22.04 LTS システムの AWS Graviton3c7g.2xlarge インスタンスでテストされています。

この例を実行するには、少なくとも4つのコアと8GBのRAMが必要です。ディスクストレージは少なくとも32GBまで設定してください。同じかそれ以上のスペックのインスタンスを使用することを推奨する。

インスタンスを起動したら、インスタンスに接続し、以下のコマンドを実行して環境を準備します。

サーバに python をインストールする：

$ sudo apt update
$ sudo apt install python-is-python3 python3-pip python3-venv -y

仮想環境を作成し、有効化します：

$ python -m venv venv
$ source venv/bin/activate

必要なpythonの依存関係をインストールします：

$ pip install --upgrade pymilvus openai requests langchain-huggingface huggingface_hub tqdm

オフラインデータのロード

コレクションの作成

ArmベースのマシンでAWS上にデプロイされたZilliz Cloudを使用して、ベクトルデータを保存し、取得します。クイックスタートするには、Zilliz Cloudに無料でアカウントを登録するだけです。

Zilliz CloudのPublic EndpointとApi keyとして uri 、token 。

from pymilvus import MilvusClient

milvus_client = MilvusClient(
    uri="<your_zilliz_public_endpoint>", token="<your_zilliz_api_key>"
)

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=384,
    metric_type="IP",  # Inner product distance
    consistency_level="Strong",  # Strong consistency level
)

デフォルトのメトリックタイプとして内積距離を使用する。距離タイプの詳細については、類似度メトリクスのページを参照してください。

データの準備

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

データの挿入

テキストを埋め込みベクトルに変換できる、シンプルだが効率的な埋め込みモデルall-MiniLM-L6-v2を用意する。

from langchain_huggingface import HuggingFaceEmbeddings

embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

テキスト行を繰り返し、埋め込みを作成し、milvusにデータを挿入する。

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

from tqdm import tqdm

data = []

text_embeddings = embedding_model.embed_documents(text_lines)

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

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

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

Arm上でLLMサービスを起動する

このセクションでは、ArmベースのCPU上でllama.cpp サービスを構築し、起動します。

Llama 3.1モデルとllama.cpp

Meta社のLlama-3.1-8Bモデルは、Llama 3.1モデルファミリーに属し、研究および商用目的で自由に使用できます。このモデルを使用する前に、Llamaのウェブサイトにアクセスし、フォームに記入してアクセスをリクエストしてください。

llama.cppはオープンソースのC/C++プロジェクトで、ローカルでもクラウドでも、様々なハードウェア上で効率的なLLM推論を可能にします。llama.cpp を使えば、Llama 3.1モデルを簡単にホストすることができます。

llama.cppをダウンロードしてビルドする

以下のコマンドを実行して、llama.cppをソースからビルドするのに必要なmake、cmake、gcc、g++、その他の必須ツールをインストールする：

$ sudo apt install make cmake -y
$ sudo apt install gcc g++ -y
$ sudo apt install build-essential -y

これでllama.cpp のビルドを始める準備ができました。

llama.cpp のソースリポジトリをクローンします：

$ git clone https://github.com/ggerganov/llama.cpp

デフォルトでは、llama.cpp は Linux と Windows でのみ CPU 用にビルドされます。実行するArm CPU用にビルドするために、余計なスイッチを用意する必要はない。

make を実行してビルドする：

$ cd llama.cpp
$ make GGML_NO_LLAMAFILE=1 -j$(nproc)

help コマンドを実行して、llama.cpp が正しくビルドされたことを確認する：

$ ./llama-cli -h

llama.cpp が正しくビルドされていれば、help オプションが表示されます。出力スニペットは次のようになります：

example usage:

    text generation:     ./llama-cli -m your_model.gguf -p "I believe the meaning of life is" -n 128

    chat (conversation): ./llama-cli -m your_model.gguf -p "You are a helpful assistant" -cnv

これでhuggingface cliを使ってモデルをダウンロードすることができます：

$ huggingface-cli download cognitivecomputations/dolphin-2.9.4-llama3.1-8b-gguf dolphin-2.9.4-llama3.1-8b-Q4_0.gguf --local-dir . --local-dir-use-symlinks False

llama.cppチームによって導入されたGGUFモデルフォーマットは、圧縮と量子化を用いて重みの精度を4ビット整数に落とし、計算量とメモリ需要を大幅に削減し、Arm CPUをLLM推論に有効にします。

モデル重みの再量子化

再量子化するには、以下を実行します。

$ ./llama-quantize --allow-requantize dolphin-2.9.4-llama3.1-8b-Q4_0.gguf dolphin-2.9.4-llama3.1-8b-Q4_0_8_8.gguf Q4_0_8_8

これは、llama-cli がSVE 256とMATMUL_INT8サポートを使用できるように再設定された重みを含む、新しいファイルdolphin-2.9.4-llama3.1-8b-Q4_0_8_8.gguf を出力します。

LLMサービスの開始

llama.cppサーバープログラムを利用し、OpenAI互換のAPI経由でリクエストを送信することができます。これにより、LLMの起動と停止を繰り返すことなく、LLMと何度もやりとりするアプリケーションを開発することができます。さらに、LLMがネットワーク経由でホストされている別のマシンからサーバーにアクセスすることもできます。

コマンドラインからサーバーを起動すると、8080番ポートをリッスンします：

$ ./llama-server -m dolphin-2.9.4-llama3.1-8b-Q4_0_8_8.gguf -n 2048 -t 64 -c 65536  --port 8080

'main: server is listening on 127.0.0.1:8080 - starting the main loop

また、起動したLLMのパラメータを調整して、サーバーのハードウェアに適合させ、理想的なパフォーマンスを得ることもできます。パラメータの詳細については、llama-server --help コマンドを参照してください。

このステップを実行するのに苦労した場合は、公式ドキュメントを参照してください。

これで、ArmベースのCPU上でLLMサービスが起動しました。次に、OpenAI SDKを使用してサービスと直接対話します。

オンラインRAG

LLMクライアントと埋め込みモデル

LLMクライアントを初期化し、エンベッディングモデルを準備します。

LLMでは、OpenAI SDKを使って、先に起動したLlamaサービスをリクエストします。実際にはローカルのllama.cppサービスなので、APIキーを使う必要はありません。

from openai import OpenAI

llm_client = OpenAI(base_url="http://localhost:8080/v1", api_key="no-key")

テスト埋め込みを生成し、その次元と最初のいくつかの要素を表示します。

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

384
[0.03061249852180481, 0.013831384479999542, -0.02084377221763134, 0.016327863559126854, -0.010231520049273968, -0.0479842908680439, -0.017313342541456223, 0.03728749603033066, 0.04588735103607178, 0.034405000507831573]

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

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

question = "How is data stored in milvus?"

コレクションから質問を検索し、意味的にトップ3にマッチするものを取得します。

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

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

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.6488019824028015
    ],
    [
        "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.5974207520484924
    ],
    [
        "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.5833579301834106
    ]
]

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

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

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

Define system and user prompts for the Language Model. This prompt is assembled with the retrieved documents from 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>
"""

LLMを使って、プロンプトに基づいたレスポンスを生成する。パラメータmodel はllama.cppサービスの冗長なパラメータなので、not-used に設定する。

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

Milvus stores data in two types: inserted data and metadata. Inserted data, including vector data, scalar data, and collection-specific schema, are stored in persistent storage as incremental log. Milvus supports multiple object storage backends such as MinIO, AWS S3, Google Cloud Storage (GCS), Azure Blob Storage, Alibaba Cloud OSS, and Tencent Cloud Object Storage (COS). Metadata are generated within Milvus and each Milvus module has its own metadata that are stored in etcd.

おめでとうございます！あなたは、Armベースのインフラストラクチャの上にRAGアプリケーションを構築しました。