OSS版EmbeddingモデルをAWSで試用した例

背景や目的


1.RAG(検索拡張生成)を中心に、本番ツールの開発研究を色々としています。

2.RAGのEmbeddingモデルはOpenAIのものに頼る事が多いのですが、性能とコストの面で有利ならOSSを使って良いのかも・・と最近思いまして、今回試してみました。

3.日本語性能が高そうなモデルを選び、それを使った検索処理をAWS上に作成します。ローカルからもクラウドからも利用できるエンドポイントを用意したく、その前段までやってみます。

使用するEmbeddingモデル


・結論として、こちらのモデルを使わせて頂きました。
 GLuCoSE v2

・日本語性能が良さそうな候補はこの他にSimCSEやMultilingual-E5等がありました。以下を参照させて頂きましたので、詳細はこちらをご覧ください。
 オープンな日本語埋め込みモデルの選択肢(NTTコミュニケーションズ株式会社杉本様)
 日本語テキスト埋め込みベンチマークJMTEBの構築(SB Intuitions李聖哲様他お二方)

実験環境でまず試す


良く使っているada-002と比べてどうなのか、軽くGoogle Colabで試してみます。

1.colabo上でセットアップ
・セル内で以下コードを実行し、必要ライブラリをインストールします。

Google Colab
###必要ライブラリインストール
pip install -U sentence-transformers

2.モデルを使うコードを用意
・モデルはHugging FaceからDLし、ベクターDB側に格納する想定の簡単な4つの文をベクター値にしてみます。
・さらに、その値の類似度判定もします。コードは以下の通りです。

Google Colab
###基本的な動作を確認
from sentence_transformers import SentenceTransformer
import torch.nn.functional as F

# Download from the 🤗 Hub
model = SentenceTransformer("pkshatech/GLuCoSE-base-ja-v2")
sentences = [
    'passage: こんにちは',
    'passage: こんにちは朝ですよ',
    'passage: さようなら',
    'passage: さようならまたあした',
]# 公式ページの情報に従い、passage:を付記。

embeddings = model.encode(sentences,convert_to_tensor=True)
print(embeddings.shape,type(embeddings))
similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
print(similarities)

・公式ページの情報をみて、ベクター値生成する元情報には、passage:を付記しました。検索時の入力文にあたる情報にはquery: を使うようです。
・Shift+Enterで実行した結果は以下の通りです。

・「こんにちは」と「こんにちは朝ですよ」は似ている、「さようなら」と「さようならまたあした」が似ている、という数値なので良し!です。次に進みます。

AWSで構築:Lambdaレイヤーの試行錯誤


本番ツールでの利用を想定し、AWSで構築したいと思います。初めに試行錯誤がありまして、上手くいかなかった部分ですが、ご参考までに少しだけ紹介します。

1.Lambda関数を作成する
・クラウドで処理をさせたく、サーバレス環境が良いのでLambdaで実現する事にしました。
・AMCの画面を使い、以下の設定値にてLambda関数を作成します。

設定項目
関数名ctcrag_trial ※お好きな名称で
ランタイムpython3.11
アーキテクチャx86_64(デフォルト)
実行ロールatd_role_for_lambda ※お手元の環境で設定
関数 URL を有効化チェック
認証タイプNONE

・画面でLambdaメニューに遷移し、トップ画面で「関数の作成」を押し、画面に従い作成します。以下は作成後の一覧画面です。ここに関数名が表示されたらOKです。

2.colaboで試したコードを移植
・先程のcolaboのコードをコード欄(lambda_function.py)にコピーし少し手直ししました。

lambda_function.py
from sentence_transformers import SentenceTransformer
import torch.nn.functional as F

# Download from the 🤗 Hub
model = SentenceTransformer("pkshatech/GLuCoSE-base-ja-v2")


def lambda_handler(event, context):
    sentences = ['passage: こんにちは','passage: こんにちは朝ですよ','passage: さようなら','passage: さようならまたあした',]

    embeddings = model.encode(sentences,convert_to_tensor=True)
    similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
    print(similarities)
    # TODO implement
    return {
        'statusCode': 200,
        'body': json.dumps('Finish!!')
    }

・実行してもライブラリが使えないので、必要ライブラリをインストールします。

3.必要ライブラリをインストール
・Lambdaの場合はレイヤーで対応するので、それをします。
・torchが大きい事を分かっていましたので、それを先に対応します。まずzipを作成しました。

Ubuntu
y-nishihara@PC999:~$ source venv/bin/activate
(venv) y-nishihara@PC999:~$ python -V
Python 3.11.0rc1
(venv) y-nishihara@PC999:~$ pip install -t ./python/ torch --extra-index-url https://download.pytorch.org/whl/cpu
(・・途中省略・・)
Installing collected packages: mpmath, typing-extensions, sympy, networkx, MarkupSafe, fsspec, filelock, jinja2, torch
Successfully installed MarkupSafe-3.0.2 filelock-3.16.1 fsspec-2024.10.0 jinja2-3.1.4 mpmath-1.3.0 networkx-3.4.2 sympy-1.13.1 torch-2.5.0+cpu typing-extensions-4.12.2
(venv) y-nishihara@PC999:~$ zip -r torch.zip ./python
(venv) y-nishihara@PC999:~$ ls -la ./torch.zip
-rw-r--r-- 1 y-nishihara y-nishihara 219375892 Oct 24 17:12 ./torch.zip

・出来たzipを画面「レイヤーの作成」で普通にアップして、まずダメでした(下図の赤字)。50MBの制限にかかります。

・今度はS3に設置して参照させる方法で試しました。

・ダメです。今度は展開後のサイズは262MB以下、という制限にかかりました。

・この後、不要部分を削除した再トライを数回やったのですが、、結果ダメでした。サイズが262MB以下になったら入りましたが、動作がダメで、、結局この方法は諦めました。

AWSで構築:コンテナ環境の作成


ここからは動作できた方法を紹介します。レイヤーでの登録は無理がありそうなので、コンテナを使う方法で構築する事にしました。

1.手元のlinux環境でdockerをインストール
・最終的にLambdaで処理させますが、まず適当なlinux環境でその準備をします。
・コンソールを開きdockerをインストールします。

Ubuntu
y-nishihara@PC999:~$ sudo apt update
(成功!メッセージ省略)
y-nishihara@PC999:~$ sudo apt install docker.io
(成功!メッセージ省略)
y-nishihara@PC999:~$ docker --version
Docker version 24.0.7, build 24.0.7-0ubuntu2~22.04.1

・とにかく入ればOKです。私の場合は上のバージョンでした。

2.Dockerfileを作成
・処理に必要な2つのファイルを作成しますが、まずDockerfileを作成します。

Ubuntu
y-nishihara@PC999:~$ mkdir my_lambda_container
y-nishihara@PC999:~$ ls
GLuCoSE-base-ja-v2  llm_work  my_lambda_container  python  rag_work  venv
y-nishihara@PC999:~$ cd my_lambda_container/
y-nishihara@PC999:~/my_lambda_container$ touch Dockerfile
y-nishihara@PC999:~/my_lambda_container$ vi Dockerfile

・作業用ディレクトリを作成し、viを開きコードを書きました。

Dockerfile
# LambdaのPythonベースイメージを使用
FROM public.ecr.aws/lambda/python:3.11

# 必要なライブラリをインストール
RUN pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
RUN pip install -U sentence-transformers
RUN pip install sentencepiece

# Lambda関数のエントリーポイント
COPY app.py ${LAMBDA_TASK_ROOT}
CMD ["app.lambda_handler"]

・Lambda上でCUDA関連でハマりたくないので、torchをcpu版にしてみました。
・sentencepieceは、途中の試行錯誤で必要だったので追加しています。

3.app.pyを作成
・続いて、app.pyを作成します。同様にvi等を開きコードを書きます。

Ubuntu
y-nishihara@PC999:~/my_lambda_container$ touch app.py
y-nishihara@PC999:~/my_lambda_container$ vi app.py
app.py
import json
from sentence_transformers import SentenceTransformer
import torch.nn.functional as F

# Download from the gitHub
model = SentenceTransformer("pkshatech/GLuCoSE-base-ja-v2")

def lambda_handler(event, context):

    sentences = [
        'passage: こんにちは',
        'passage: こんにちは朝ですよ',
        'passage: さようなら',
        'passage: さようならまたあした',
    ]

    embeddings = model.encode(sentences,convert_to_tensor=True)
    print(embeddings.shape,type(embeddings))

    # Get the similarity scores for the embeddings
    similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
    print(similarities)

    similarity_list = similarities.tolist()  # JSON形式に変換可能なリストに変換
    return {
        'statusCode': 200,
        'body': json.dumps({'RESULTS': similarity_list})
    }

・colaboで試したコードをベースに、Lambda関数にマッチするように修正しました。
・最後にlsして作成済の2つのファイルあればOKです。

Ubuntu
y-nishihara@PC999:~/my_lambda_container$ ls -la
total 24
drwxr-xr-x  2 y-nishihara y-nishihara 4096 Nov 18 10:52 .
drwxr-x--- 16 y-nishihara y-nishihara 4096 Dec  6 17:35 ..
-rw-r--r--  1 y-nishihara y-nishihara  390 Nov  1 11:18 Dockerfile
-rw-r--r--  1 y-nishihara y-nishihara  900 Nov 15 17:56 app.py

AWSで構築:コンテナイメージ作成とECR登録


手元のlinux環境でコンテナイメージを作成し、AWS ECRへプッシュします。

1.コンテナイメージを作成
・前の章で作成した2つのファイルを使いコマンド1つで作成できます。

Ubuntu
y-nishihara@PC999:~/my_lambda_container$ sudo docker build -t lambda-container-test1 .
(・・途中省略・・)
Successfully built e0de0baa78c2
Successfully tagged my-lambda-container:latest
y-nishihara@PC999:~/my_lambda_container$ sudo docker images
REPOSITORY                     TAG       IMAGE ID       CREATED              SIZE
lambda-container-test1            latest    e0de0baa78c2   About a minute ago   9.56GB
public.ecr.aws/lambda/python   3.11      60f53643b1d9   11 days ago          692MB

・時間は初回は約11分かかりました。2回目以降は速いです。

2.Amazon Elastic Container Registry (ECR) でリポジトリ作成
・作成したコンテナイメージをプッシュする先のリポジトリをまず作成します。
・Amazon ECRのメニューを開き、Private registryの下のrepositoriesを押します。

・その画面で「リポジトリを作成」を押すと、以下作成画面が表示されます。

・ここでは、以下の値を入力し「作成」を押します。

設定項目
リポジトリ名lambda-container-test1 ※お好きな名称で
イメージタグのミュータビリティMutable
暗号化設定AES-256
その他全てデフォルト

・以下画面が出たら作成は成功です。画面内のlambda-container-test1を押します。

・イメージはまだありませんね。ここで「プッシュコマンドを表示」を押します。

・ここに書かれた1.~4.をクライアント側で実行すれば良さそうです。

3.クライアントからECRへプッシュ
・一旦クライアントに戻り、上の1.3.4.をやります。(2.は既に実施済)
・私は手元のLinux環境でAWS CLIを入れて実行しました。以下はそのログです。

Ubuntu
y-nishihara@PC999:~$ sudo apt update  
(・・途中省略・・)
56 packages can be upgraded. Run 'apt list --upgradable' to see them.
y-nishihara@PC999:~$ sudo apt install unzip curl -y
(・・途中省略・・)
0 upgraded, 0 newly installed, 0 to remove and 56 not upgraded.
y-nishihara@PC999:~$ curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100 63.4M  100 63.4M    0     0  10.3M      0  0:00:06  0:00:06 --:--:-- 10.6M
y-nishihara@PC999:~$unzip awscliv2.zip
(・・途中省略・・)
y-nishihara@PC999:~$ sudo ./aws/install
[sudo] password for y-nishihara:
You can now run: /usr/local/bin/aws --version
y-nishihara@PC999:~$ aws --version
aws-cli/2.18.15 Python/3.12.6 Linux/5.15.153.1-microsoft-standard-WSL2 exe/x86_64.ubuntu.22

・上の操作でAWS CLIは入りました。次はaws configureをします。
・そこでKey IDとAccess Keyが必要となるので、先にAMCで作成しておきます。
・AMCのIAMメニューでユーザーを開き「ユーザーの作成」を押します。

・各入力画面では、下表の値を入力し、最後に「ユーザーの作成」を押します。

設定項目
ユーザー名lambda-ecr-user ※お好きな名称で
許可のオプションポリシーを直接アタッチする
許可ポリシーAmazonEC2ContainerRegistryFullAccess
AWSLambda_FullAccess
その他全てデフォルト

・作成したユーザーを選び、その画面で「アクセスキーを作成」を押します。

・ここでは、コマンドラインインターフェイス (CLI)を選び、「次へ」を押します。

・説明タグの設定では、そのまま「アクセスキーを作成」を押します。

・Key ID(アクセスキー)とAccess Key(シークレットアクセスキー)を入手しました。
・クライアントコンソールに戻り、それらを使いaws configureをし、最後にpushします。

Ubuntu
y-nishihara@PC999:~$ aws configure
AWS Access Key ID [None]: ***上で得たアクセスキー
AWS Secret Access Key [None]: ***上で得たシークレットアクセスキー
Default region name [None]: ap-northeast-1
Default output format [None]:
y-nishihara@PC999:~$ aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin ***以降省略 AMC画面で表示された認証コマンドのコピーでOK
permission denied while trying to connect to the Docker daemon socket at unix:///var/run/docker.sock: Post "http://%2Fvar%2Frun%2Fdocker.sock/v1.24/auth": dial unix /var/run/docker.sock: connect: permission denied
y-nishihara@PC999:~$ sudo usermod -aG docker $USER
y-nishihara@PC999:~$ exit ★1回閉じて、再度コンソールを開く
y-nishihara@PC999:~$ aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin ***以降省略 AMC画面で表示された認証コマンドのコピーでOK
https://docs.docker.com/engine/reference/commandline/login/#credentials-store
Login Succeeded
y-nishihara@PC999:~$ docker tag lambda-container-test1:latest ***以降省略 AMC画面で表示されたタグ付けコマンドのコピーでOK
y-nishihara@PC999:~$ docker push ***以降省略 AMC画面で表示されたpushコマンドのコピーでOK
The push refers to repository [***.dkr.ecr.ap-northeast-1.amazonaws.com/lambda-container-test1]
***: Pushed
(・・途中省略・・)
***: Pushed
latest: digest: sha256:d4**+3d size: 2000

・途中のpermission deniedでは、Dockerデーモンへのアクセス権を付加し対処しています。
・認証成功後にタグ付けしています(上で記載したAMC画面のプッシュコマンド3.にて)。
・最後のプッシュは(初回は)25分位かかりました。2回目以降は速いです。
・これでECRへのプッシュが出来ました!

AWSで構築:コンテナをLambdaにセット


コンテナイメージをECRへプッシュ出来たので、それをLambdaで使えるようにします。

・上の方で作成したLambda関数は使えないので、改めてコンテナ版の関数を作成します。
・AMCのLambdaメニューで「関数の作成」を押します。

・以下の画面「関数の作成」では、下表内の値を入力し「関数の作成」を押します。

設定項目
関数作成方法コンテナイメージ ※重要です
関数名ctcrag_trial2 ※お好きな名称で
コンテナイメージURI ***.dkr.ecr.ap-northeast-1.amazonaws.com/lambda-container-test1@sha256:d4***3d
※上でpushしたイメージです。@sha256:・・も必要です
アーキテクチャx86_64
実行ロールatd_role_for_lambda ※お手元の環境に合わせて下さい
※別途IAM画面で”AmazonEC2ContainerRegistryFullAccess”をアタッチします
その他全てデフォルト

・コンテナイメージURLは、「イメージを参照」を押してセットすると楽です。


・「関数の作成」押下後、3分位で作成完了の表示がされます。

・これでコンテナ版Lambda関数が作成され、先ほど作成したコンテナイメージも紐付きました。

動作確認1


コンテナイメージをLambda関数で使ってみます。

1.とにかくアクセス
・AMCでLambdaメニューに入り、先ほど作成した関数を選び、テストタブを押します。

・適当なイベント(図ではtest1)を作成しておき、右上の「テスト」を押します。

・はじめはエラーが出て、動作するまで少し試行錯誤しました。

2.追加の設定
・ログをみて、以下の3つに注目し、

1) ・・・ Task timed out after 3.03 seconds
2) There was a problem when trying to write in your cache folder (/home/sbx_user1051/.cache/huggingface/hub). You should set the environment variable TRANSFORMERS_CACHE to a writable directory.
3) /var/lang/lib/python3.11/site-packages/transformers/utils/hub.py:128: FutureWarning: Using TRANSFORMERS_CACHE is deprecated and will be removed in v5 of Transformers. Use HF_HOME instead.

・最終的には以下の5つをLambdaの画面で追加登録しました。

設定項目設定場所
メモリ2048MB<本Lambda関数>/設定/一般設定
エフェメラルストレージ1024MB<本Lambda関数>/設定/一般設定
タイムアウト3分<本Lambda関数>/設定/一般設定
TRANSFORMERS_CACHE /tmp<本Lambda関数>/設定/環境変数
HF_HOME/tmp<本Lambda関数>/設定/環境変数

3.再度アクセス
・同じ方法でテスト実行をし、正常に処理がされました!機能的にはこれでOKです。

・動作確認時に、始めのアクセスが遅い、という事が気になりました。
・色々調べてみると、それはコールドスタートとウォームスタートの違いだと理解できました。しばらく使っていないと遅く、その逆は速い、というものです。
・コールドスタート時は約20秒、ウォームスタート時(直ぐ2度目を実行した時)は1秒位、の処理時間でした。
・今回はアクセス間隔が5分を超えるとコールドスタートになっていました。ウォームスタート率を高くする事がテーマですね。構築時には考慮したいと思います。

AWSで構築:モデルをS3に設置


機能的には前章で成功していますが、モデルがHugging Faceからダウンロードされる点が気になり、S3に設置して動作するか試しました。蛇足かな・・

1.モデルのデータをファイル保存
・S3に置くファイルを作ります。colaboでやりました。

Google Colab
# コンテナ版Lambda関数作成で実行したpip installをここでも実行
!pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
!pip install -U sentence-transformers
!pip install sentencepiece
# モデルをダウンロードし、model.saveで保存
import json
from sentence_transformers import SentenceTransformer
import torch.nn.functional as F
model = SentenceTransformer("pkshatech/GLuCoSE-base-ja-v2")
model.save("/content/GLuCoSE-base-ja-v2-model")

・作成された/content/GLuCoSE-base-ja-v2-modelは、colaboのメニューで操作可能です。まずはお手元の適当な場所にコピーしておきます。

2.S3へアップロード
・S3では適当なバケットを作成し、そこにディレクトリごとアップロードします。

S3設定項目
バケット名ctcrag-trial ※お好きな名称で

・以下はアップロード後の画面です。(S3の操作方法は公式ページ等をご参照下さい。)

・フォルダの中には色々と入っています。

3.処理内容を微修正
・コンテナのファイルを修正します。Dockerfileは変更不要なのでapp.pyのみ修正します。

app.py
import json
from sentence_transformers import SentenceTransformer
import torch.nn.functional as F
import boto3 #added for S3
import os #added for S3

bucket_name = 'ctcrag-trial'  # S3バケット
model_prefix = 'GLuCoSE-base-ja-v2-model'

def lambda_handler(event, context):
	# Lambdaの一時フォルダにダウンロード
	local_model_dir = '/tmp/model/'
	snapshot_dir = os.path.join(local_model_dir,model_prefix)
	os.makedirs(local_model_dir, exist_ok=True)
	# S3からモデル全体をコピー
	s3_client = boto3.client('s3')
	if not os.path.exists(snapshot_dir):
		for obj in s3_client.list_objects_v2(Bucket=bucket_name, Prefix=model_prefix)['Contents']:
			download_path = os.path.join(local_model_dir, obj['Key'])
			os.makedirs(os.path.dirname(download_path), exist_ok=True)
			s3_client.download_file(bucket_name, obj['Key'], download_path)
	# モデルのロード
	model = SentenceTransformer(snapshot_dir)
	#ベクター値作成
	sentences = ['passage: こんにちは','passage: こんにちは朝ですよ','passage: さようなら','passage: さようならまたあした',]
	embeddings = model.encode(sentences,convert_to_tensor=True)
	#類似度計算
	similarities = F.cosine_similarity(embeddings.unsqueeze(0), embeddings.unsqueeze(1), dim=2)
	print(similarities)
	similarity_list = similarities.tolist()  # JSON形式に変換可能なリストに変換
	return {
		'statusCode': 200,
		'body': json.dumps({'RESULTS': similarity_list})
	}

4.コンテナイメージ再作成&ECRプッシュ&Lambda関数へ紐づけ
・これらを実行したら完了です。
・手順はこれまでに記載した方法と同様です。そちらを参照して下さい。

動作確認2(モデルS3設置版)


S3に設置したモデルを使った処理を試してみます。

1.Lambda関数のテスト実行
・Lambda関数は実行可能な状態なので、画面上でテスト実行します。

・はい。成功しました。

2.処理時間について
・S3版では初回実行時には、約1分30秒かかりました。
・その後、コールドスタート時とウォームスタート時をそれぞれ計測をしたら、前者は約20秒で後者は1秒以内、という結果でした。Hugging Face版と同じです。(上の初回だけ何等かのオーバーヘッドがあるようです。)
・動作確認1でも書きましたが、やはりウォームスタート率を高める事が大事になりそうです。

システム構成


それぞれの構成を図示しておきます。

まとめ


1.OSSのEmbeddingモデルをAWS上に設置して使う事が出来ました。

2.今回はモデル自体をHugging Faceから使う形と、S3等に置く形を両方試しそれぞれ機能しましたが、どちらが良いか本番構築時に考えないといけないですね。(どっちが良いのかなぁ)

3.モデルは無償でも、Lambdaの費用がかかる事に注目しました。本編ではあまり触れませんでしたが、コールドスタートを避けるためにProvisioned Concurrencyを使うとそこそこお金がかかりますね。(最低でも月3000円位でしょうか)

4.今関わっている案件で、今回の方式を試してみます。RAGの構築は今後もしばらくの間は大事なテーマになりそうなので、どんどんトライします。その内容はまた報告致します。

お読み頂いてありがとうございました。怪しい部分があるかもしれませんので、何かお気づきの点があれば、ご指摘頂くと大変助かります。
下の「ご連絡フォーム」の「・この記事に関して」から入力可能です。

ご連絡フォーム


フィードバックを是非お願いします。
本記事の方法での問題点や、よりよい方法のアイデアを頂けると大変助かります。

この記事に関して

その他のご連絡

DevAIsをもっと見る

今すぐ購読し、続きを読んで、すべてのアーカイブにアクセスしましょう。

続きを読む