Amazon Titan Text Embedding V2 を試してみた

つくぼし

2024.05.06 08:21

こんにちは、つくぼし(tsukuboshi0755)です！

Amazon Titan Text EmbeddingsのV2が、Amazon Bedrockの新規埋め込みモデルとして利用可能になりました。

今回はこちらのAmazon Titan Text Embeddings V2を試してみたので紹介します！

G1とV2の比較

G1とV2の比較表は以下の通りです。

特徴	Titan Text Embeddings G1	Titan Text Embeddings V2
最大トークン数	8k	8k
ベクトルサイズ	1532	256/512/1024
対応言語数	25以上	100以上
正規化のサポート	なし	あり
100万トークンあたりの価格	0.10USD	0.02USD

V2はG1に比べて以下のようなメリットがあります。

より小さいベクトルサイズを選択する事で、データベースから文書抽出を検索して取得するための待ち時間を短縮可能
ベクトルの類似性を測定する際の精度の向上に役立つ、ベクトルの正規化をサポート
より幅広い言語がサポートされ、多言語の文書を埋め込む際に役立つ
トークンあたりの価格が低下し、さらにコストを削減可能

なお2024/5/6現在、Amazon Titan Text Embeddings V2はナレッジベースの埋め込みモデルとしては使用できないためご注意ください。

試してみた

今回はPythonのboto3を用いてTitan Text V2のAPIを呼び出し、2つの文を入力した上で、2つの文章の類似性を評価するコサイン類似度を計算してみます。

コサイン類似度については、以下のブログが詳しいので別途ご参照ください。

パラメータ

V2ではG1と比較して、APIで指定可能なリクエストボディのパラメータが増えているため、以下で紹介します。

パラメータ	指定可能な値	デフォルト値	用途
inputText	(文字列)	-	埋め込みベクトルに変換する入力テキスト。
normalize	true/false	true	埋め込みベクトルを正規化するかどうかを示すフラグ。 RAGで使用する場合はtrueでの設定が推奨される。
dimensions	256/512/1024	1024	埋め込みベクトルが持つ次元の数。待ち時間を短縮したい場合は小さく、逆に精度を高めたい場合は大きい値を指定すると良い。

特にV2ではベクトルサイズ(dimensions)を柔軟に選択できるようになったため、用途に応じて適切な値を選択すると良さそうです。

使用ライブラリ

今回使用するライブラリは、boto3のみです。

boto3：1.34.61

コサイン類似度を計算する際、一般的には以下のブログのようにNumPyやSciPyという外部ライブラリを別途導入し利用する事が多いです。

今回はあえて外部ライブラリを使わず、boto3以外はPythonの標準ライブラリのみでコサイン類似度を計算します。

boto3がプリインストールされているCloudShell等の環境であれば追加の外部ライブラリなしで実行できるので、ぜひそちらでお試しください。

スクリプト内容

Pythonのスクリプトを以下に示します。

スクリプト実行時の引数に比較対象となる2つの文を指定して実行すると、コサイン類似度が出力されます。

また今回はオレゴンジージョン(us-west-2)のBedrock呼び出しを想定しています。

import boto3
import json
import math
import sys

# メイン処理
def main():
    # 引数から比較対象となる2つの文を取得
    sentence1 = sys.argv[1]
    sentence2 = sys.argv[2]

    # 2つの文の埋め込みベクトルを計算
    vector1 = calculate_vector(sentence1)
    vector2 = calculate_vector(sentence2)

    # 2つの埋め込みベクトルのコサイン類似度を計算して出力
    similarity = cosine_similarity(vector1, vector2)
    print(similarity)

# テキストを入力として受け取り、埋め込みモデルを呼び出してベクトルを計算する関数
def calculate_vector(sentence):
    # Bedrock Runtimeクライアントを作成
    bedrock_runtime = boto3.client(
        service_name="bedrock-runtime", region_name="us-west-2"
    )

    # ベクトルサイズを256/512/1024のいずれかで指定
    dimensions = 256

    # リクエストボディを作成
    body = json.dumps(
        {
            "inputText": sentence,
            "dimensions": dimensions,
        }
    )

    # Bedrockの埋め込みモデルIDを指定
    modelId = "amazon.titan-embed-text-v2:0"

    # モデルを呼び出してレスポンスを取得
    response = bedrock_runtime.invoke_model(
        body=body,
        modelId=modelId,
        accept="*/*",
        contentType="application/json",
    )

    # レスポンスボディを取得
    response_body = json.loads(response.get("body").read())

    # 埋め込みベクトルを取得して返す
    embedding = response_body.get("embedding")
    return embedding

# ベクトルの内積を計算する関数
def dot_product(vecA, vecB):
    return sum(x * y for x, y in zip(vecA, vecB))

# ベクトルのユークリッドノルム(大きさ）を計算する関数
def magnitude(vec):
    return math.hypot(*vec)

# コサイン類似度を計算する関数
def cosine_similarity(vec1, vec2):
    return dot_product(vec1, vec2) / (magnitude(vec1) * magnitude(vec2))

if __name__ == "__main__":
    main()

実行結果

スクリプト実行前に、対象リージョンにおけるBedrockコンソールのモデルアクセスより、Amazon Titan Text Embeddings V2を有効化する必要があるためご注意ください。

有効化後、CloudShell上でmain.pyを実行した結果を以下に示します。

$ python main.py "東京の醤油ラーメン" "東京の味噌ラーメン"
0.8402169065862879

$ python main.py "東京の醤油ラーメン" "札幌の味噌ラーメン"
0.49552747137073716

$ python main.py "東京の味噌ラーメン" "札幌の味噌ラーメン"
0.5408871354010766

作成したスクリプトを実行すると、上記の通り問題なくコサイン類似度が計算できているようです。

例えば「東京の醤油ラーメン」、「東京の味噌ラーメン」、「札幌の味噌ラーメン」の中では、「東京の醤油ラーメン」と「東京の味噌ラーメン」のコサイン類似度が最も高くなっていることが確認できます。