NeMo Retriever を使って Embedding モデルを簡単デプロイ

Reading Time: 3 minutes

NeMo Retriever は、エンタープライズデータの意味検索 (Semantic Search) を可能にし、検索拡張 (Retrieval Augmentation) を使用して高精度の応答を提供するマイクロサービスのコレクションです。本記事では、NeMo Retriever に含まれている NeMo Retriever Embedding Microservice を使用して Embedding モデルを簡単にデプロイする方法を説明します。

NeMo Retriever とは

NeMo Retriever は、カスタム生成 AI モデルをさまざまなビジネスデータとシームレスに接続し、高精度の応答を提供するマイクロサービスのコレクションです。NeMo Retriever には、最小限のレイテンシー、最高のスループット、最大限のデータプライバシーを備えた情報検索機能が含まれており、組織はデータをより有効に活用し、ビジネス上の洞察をリアルタイムに生成できます。

NeMo Retriever は現在、Early Access のフェーズで NeMo Microservices Early Acces から申し込みいただき、承認後にアクセスすることができます。承認後は以下のサービスを利用できます。また、この記事でご案内している、ハードウェア要件、対応モデル一覧、モデルのダウンロード方法などの NeMo Microservices 関連ドキュメントにアクセスするには、早期アクセスへのご登録が必要です。

NeMo Retriever Embedding Microservice
- CUDA、TensorRT、Triton Inference Server を組み込んだ NVIDIA ソフトウェアプラットフォーム上に構築されており、OpenAI 互換の API を介して高速化された Embedding モデルの推論サービングを提供する Microservice。
NeMo Retriever Reranking Microservice
- ユーザーの入力やコンテキスト、またはアプリケーション固有の基準に基づいて検索された文書を並び替え、入力と検索結果との関連性をより高めるための Microservice 。
NeMo Retriever Microservice
- ユーザーデータのインデックス作成とクエリを行うための単一の API を提供する Microservice。あらかじめ組み込まれているパイプライン機能には、チャンキング、Embedding、ベクトルデータベースなどが含まれています。

Retrieval Augmented Generation とは

Retrieval Augmented Generation (RAG) とは、図1. RAG のアーキテクチャ例にあるように、企業のナレッジベースなどから取得した情報を用いて、LLM の精度と信頼性を向上させるアプローチです。

RAG のプロセスは大まかに以下の通りです。

企業のナレッジベースの取り込み。これにはテキスト、PDF、HTML などが含まれ、後続のプロセスで検索して利用できるよう Embedding モデルで文の意味を考慮したベクトル表現へ変換され、ベクトルデータベースへ格納されます。
ユーザーからのクエリと応答の生成。ユーザーからのクエリは、 Embedding モデルでベクトルへ変換され、ベクトルデータベース内から意味的な類似性を基に関連データが検索されます。それらはユーザーからのクエリとともにコンテキストとして追加され、LLM は応答を生成します。

RAG は、LLM に更新可能な膨大な知識を提供し、リアルタイムなイベントや特定の知識ドメインなどへ効果的に対処できるようにします。また、RAG は参照したデータソースを応答に含めることもできるため、ユーザーは LLM が生成した応答と引用されたソースを組み合わせて情報を活用することが可能になります。

Embedding モデルは RAG のワークフロー内で、頻繁に更新されるナレッジベースの取り込みやユーザーからのクエリの度に呼び出されることになるため、高い推論処理パフォーマンスとスケーラビリティが求められます。

NeMo Retriever Embedding Microservice チュートリアル

本記事では、Hugging Face Model Hub から Embedding モデルをダウンロードして、NeMo Retriever Embedding Microservice (NREM) を使用してデプロイします。

本チュートリアルでの手順は以下の通りです。

NREM を起動するための事前準備
Hugging Face Model Hub からモデルをダウンロード
NREM の起動
推論の実行
LangChain から使用する方法
(オプション): 日本語 Embedding モデルを使用する例

また、今回のチュートリアルの検証環境は以下の条件で行っています。

ハードウェア
- GPU: 1x NVIDIA L4 (driver version: 535.54.03)
- CPU: Dual Intel(R) Xeon(R) Gold 5318Y (24 コア, 2.1 GHz)
- システムメモリ: 64 GB
ソフトウェア
- OS: Ubuntu 22.04.4 LTS
- Docker: 26.0.2
- nvidia-container-toolkit: 1.15.0
- Git LFS: 3.0.2
- python 3.10.12
- langchain-0.1.16
- langchain-community-0.0.34
- 使用した Docker イメージ: nvcr.io/ohlfw0olaadg/ea-participants/nemo-retriever-embedding-microservice:24.02

NREM の HW/SW 要件はこちらに記載があります。

NREM を起動するための事前準備

GPU が搭載されたマシンを準備します。
NGC から最新の NeMo Retriever Embedding Microservice コンテナーを入手します (このチュートリアルでは、nvcr.io/ohlfw0olaadg/ea-participants/nemo-retriever-embedding-microservice:24.02 を使用しています)。
以下のコマンドで作業用のディレクトリを作成し、移動します。

mkdir nrem-example
cd nrem-example

Hugging Face Model Hub からモデルをダウンロード

NREM では現在、以下のモデルをサポートしています。

NV-Embed-QA
e5-large-v2 を含む E5 の派生
gte-large を含む GTE の派生
gtr-t5-base を含む GTR の派生

最新のサポートモデルの一覧はこちらで確認してください。このチュートリアルでは、日本語を含む多言語に対応したモデルである intfloat/multilingual-e5-large を使用します。以下のスクリプトで Hugging Face の Model Hub からモデルをダウンロードします。他のモデルのダウンロードの方法はこちらを参照してください。

git lfs install
git clone https://huggingface.co/intfloat/multilingual-e5-large

NREM の起動

モデルのダウンロードが完了したら次に NREM を起動します。以下のスクリプトを実行すると先ほどダウンロードしたモデルが TensorRT 形式に変換された後に Triton Inference Server 上でサービングされます。

export MODEL_ID=e5
export MODEL_NAME=multilingual-e5-large

export HOSTNAME="0.0.0.0"
export SERVICE_PORT=8080
export CONTAINER_NAME="embedding-ms"

docker run \
  -it \
  --rm \
  --name $CONTAINER_NAME \
  --gpus=1 \
  --shm-size=8G \
  -v $(pwd):/model-checkpoint-path \
  -p $SERVICE_PORT:$SERVICE_PORT \
  nvcr.io/ohlfw0olaadg/ea-participants/nemo-retriever-embedding-microservice:24.02 \
  bin/web -p $SERVICE_PORT -c /model-checkpoint-path/$MODEL_NAME -g model_config_templates/${MODEL_ID}_template.yaml

以下のようなログが標準出力に出て、サーバーが起動すれば成功です。

I0228 12:06:51.173629 352 grpc_server.cc:2513] Started GRPCInferenceService at 0.0.0.0:8001
I0228 12:06:51.173953 352 http_server.cc:4497] Started HTTPService at 0.0.0.0:8000
I0228 12:06:51.215066 352 http_server.cc:270] Started Metrics Service at 0.0.0.0:8002

先ほどのスクリプトでは、実行のたびにモデルを TensorRT 形式へと変換する処理が行われるため、起動までに数分ほど時間がかかってしまいます。以下のスクリプトを実行することで、TensorRT 形式のモデルをローカル環境へエクスポートすることができます。

export MODEL_ID=e5
export MODEL_NAME=multilingual-e5-large

export HOSTNAME="0.0.0.0"
export SERVICE_PORT=8080
export CONTAINER_NAME="embedding-ms"

docker run \
  -it \
  --name generate-model-store \
  --gpus=1 \
  --shm-size=8G \
  -v $(pwd):/model-checkpoint-path \
  -p $SERVICE_PORT:$SERVICE_PORT \
  nvcr.io/ohlfw0olaadg/ea-participants/nemo-retriever-embedding-microservice:24.02 \
  model_repo_generator model_config_templates/${MODEL_ID}_template.yaml -c /model-checkpoint-path/$MODEL_NAME

docker cp generate-model-store:/model-store converted_model_repository

実行が無事完了すると、ローカル環境に作成された converted_model_repository というディレクトリの中にモデルがエクスポートされていることが確認できます。

次に、以下はすでに TensorRT 形式のモデルがエクスポートされている場合に実行するスクリプトです。変換処理がスキップされるため、コンテナーの起動は数秒ほどで完了します。

export HOSTNAME="0.0.0.0"
export SERVICE_PORT=8080
export CONTAINER_NAME="embedding-ms"

docker run \
  -it \
  --rm \
  --name $CONTAINER_NAME \
  --gpus=1 \
  --shm-size=8G \
  -v $(pwd)/converted_model_repository:/model-store \
  -p $SERVICE_PORT:$SERVICE_PORT \
  nvcr.io/ohlfw0olaadg/ea-participants/nemo-retriever-embedding-microservice:24.02 \
  bin/web -m /model-store -p $SERVICE_PORT

推論の実行

先ほどサービングされたコンテナーにクエリを投げたい場合は以下のように実行します。API は OpenAI の API と互換性があり、クエリ実行後に埋め込みベクトルを含むレスポンスが返ってきます。API の仕様についてはこちらを参照してください。

curl -X "POST" \
  "http://${HOSTNAME}:${SERVICE_PORT}/v1/embeddings" \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
    "input": ["こんにちは、世界！"],
    "model": "'"$MODEL_ID"'",
    "input_type": "query"
}'

Python からクエリを投げたい場合は以下のように実行します。

import requests
import json

HOSTNAME = "0.0.0.0"
SERVICE_PORT = 8080
model_id = "e5"

url = f"http://{HOSTNAME}:{SERVICE_PORT}/v1/embeddings"

payload = json.dumps({
  "input": [
    "こんにちは、世界！"
  ],
  "model": model_id,
  "input_type": "query"
})

headers = {
  'Content-Type': 'application/json'
}

response = requests.request("POST", url, headers=headers, data=payload)
print(response.text)

LangChain から使用する方法

NREM は LLM アプリを構築するフレームワークである LangChain に統合されており、フレームワーク上から簡単に呼び出すことができます。

from langchain_community.embeddings import NeMoEmbeddings

batch_size = 16
model = "e5"
api_endpoint_url = "http://localhost:8080/v1/embeddings"

embedding_model = NeMoEmbeddings(
    batch_size=batch_size,
    model=model,
    api_endpoint_url=api_endpoint_url
    )

# Checking if endpoint is live:
# http://localhost:8080/v1/embeddings

ret = embedding_model.embed_query("こんにちは、世界！")
print("特徴量の次元数:", len(ret))
print("特徴量:", ret)

(オプション): 日本語 Embedding モデルを使用する例

最後に NREM のカスタマイズの例として日本語の Embedding モデルを使用する方法を説明します。日本語の言語モデルでは、Mecab に代表される日本語のトークナイザーが使用されることがあります。NREM には日本語トークナイザーのライブラリが含まれていないため、ここでは fugashi という Mecab のラッパーライブラリを追加するために下記のような Dockerfile を作成します。

FROM nvcr.io/ohlfw0olaadg/ea-participants/nemo-retriever-embedding-microservice:24.02

RUN pip3 install -U fugashi[unidic-lite]

次に下記スクリプトで Dockerfile から新たに Docker イメージを構築します。

docker build -t nvcr.io/ohlfw0olaadg/ea-participants/nemo-retriever-embedding-microservice:24.02-fugashi .

この例では、日本語の Embedding モデルである cl-nagoya/sup-simcse-ja-large を例として使用するため、Hugging Face の Model Hub からモデルをダウンロードします。

git clone https://huggingface.co/cl-nagoya/sup-simcse-ja-large

次に、NREM を起動する前に下記のテンプレートファイルを simcse-ja_template.yaml という名前で保存しておきます。multilingual-e5-large の例では、このテンプレートファイルはコンテナー内にあらかじめ保存されたものを使用していました。しかし、このコンテナー内には、sup-simcse-ja-large のテンプレートファイルは用意されていないため、事前に作成して用意しておく必要があります。

このテンプレートファイルは使用するモデルによって異なる場合があり、かつコンテナーがアップデートされた際に内容が変更される可能性がある点に注意してください。

triton_store_path: "/model-store/"
models:
  - name: "simcse-ja"
    checkpoint_path: "cl-nagoya/sup-simcse-ja-large"
    triton_config:
      trt:
        num_instances: 1
        max_shapes: [128, 512]
        max_queue_delay_microseconds: 100
        dtype: float16
      pipeline:
        num_instances: 1

次に NREM を起動します。multilingual-e5-large の例との変更点は、起動する Docker イメージが変更されている点とテンプレートファイルのパスが変更されている点になります。

export MODEL_ID=simcse-ja
export MODEL_NAME=sup-simcse-ja-large

export HOSTNAME="0.0.0.0"
export SERVICE_PORT=8080
export CONTAINER_NAME="embedding-ms"

mkdir converted_model_repository

docker run \
  -it \
  --rm \
  --name $CONTAINER_NAME \
  --gpus "device=1" \
  --shm-size=8G \
  -v $(pwd):/model-checkpoint-path \
  -p $SERVICE_PORT:$SERVICE_PORT \
  nvcr.io/ohlfw0olaadg/ea-participants/nemo-retriever-embedding-microservice:24.02-fugashi \
  bin/web -p $SERVICE_PORT -c /model-checkpoint-path/$MODEL_NAME -g /model-checkpoint-path/${MODEL_ID}_template.yaml

起動が完了したら、サービングされたコンテナーへクエリを投げてみます。multilingual-e5-large の例との変更点は、このモデルでは必要としない、”input_type”: “query” という入力が省略されている点になります。

curl -X "POST" \
  "http://${HOSTNAME}:${SERVICE_PORT}/v1/embeddings" \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
    "input": ["こんにちは、世界！"],
    "model": "'"$MODEL_ID"'"
}'

まとめ

本記事では、NeMo Retriever Embedding Microservice を使用した Embedding モデルのデプロイ方法を紹介しました。NeMo Retriever Embedding Microservice を使用して LLM のワークロード構築が加速すると嬉しいです。

NeMo Retriever を使って Embedding モデルを簡単デプロイ

NeMo Retriever とは

Retrieval Augmented Generation とは

NeMo Retriever Embedding Microservice チュートリアル

NREM を起動するための事前準備

Hugging Face Model Hub からモデルをダウンロード

NREM の起動

推論の実行

LangChain から使用する方法

(オプション): 日本語 Embedding モデルを使用する例

まとめ

関連情報

Tags

About the Authors

NeMo Retriever を使って Embedding モデルを簡単デプロイ

NeMo Retriever とは

Retrieval Augmented Generation とは

NeMo Retriever Embedding Microservice チュートリアル

NREM を起動するための事前準備

Hugging Face Model Hub からモデルをダウンロード

NREM の起動

推論の実行

LangChain から使用する方法

(オプション): 日本語 Embedding モデルを使用する例

まとめ

関連情報

Tags

About the Authors

関連記事

本番環境への AI モデルのデプロイを効率化する NVIDIA NIM

NVIDIA NeMo マイクロサービスでカスタム生成 AI 開発が簡単に

LLM テクニックの習得: 推論の最適化

LLM テクニックの習得: トレーニング

LLM テクニックの習得: LLMOps