Amazon BedrockでテキストのEmbeddingsを取得してみた

こんにちは。CX事業本部Delivery部のakkyです。

Amazon Bedrockにはチャット、テキスト生成、画像生成などさまざまな生成AIのモデルが用意されていて、フルマネージドで使えます。

その中でも、Titan Embeddings Generation 1 (G1)というモデルを使うと、テキストのEmbeddingsを取得することができます。これは文章の意味をベクトルに変換するもので、テキストの類似度比較などに使用できます。 多言語対応ということで日本語の文章も使えるようです。

Titan Embeddingsモデルでは、入力された文章を1536次元のベクトル（1536個の数値の配列）に変換できます。

利用料金は1kトークン当たり$0.0001です。

Pythonコードによる取得

Amazon Bedrockではテキスト生成やチャットボットや画像生成に関してはコンソールからお手軽に試すことができますが、Embeddingsは結果がベクトルという性質からAPIからアクセスする必要があります。

今回はPythonから使ってみました。

使用ライブラリ

boto3の最新版とnumpyが必要です

• boto3 1.28.57
• numpy 1.26.0

ソースコード

AWSから提供されているAmazon Bedrock Workshopを参考にしました。Jupyter Notebook用ですが、通常のPythonコードの参考にもなります。

Embeddingは次のノートブックに出てきます。

https://github.com/aws-samples/amazon-bedrock-workshop/blob/main/00_Intro/bedrock_boto3_setup.ipynb

ベクトルの比較にはコサイン類似度を使用します。こちらの記事を参考にさせていただきました。

import boto3
import json
import numpy as np

bedrock_client = boto3.client('bedrock-runtime', region_name="us-west-2")

def bedrock_embedding(input_str: str):
    bedrock_body = {
        "inputText": input_str
    }
    body_bytes = json.dumps(bedrock_body).encode('utf-8')
    response = bedrock_client.invoke_model(
        accept="*/*",
        body=body_bytes,
        contentType="application/json",
        modelId="amazon.titan-embed-text-v1",
    )
    response_body = json.loads(response.get("body").read())

    #print(response_body.get("inputTextTokenCount"))
    embedding = response_body.get("embedding")
    return embedding

def cosine_similarity(v1, v2):
    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))


a = bedrock_embedding("埋め込みの実験のためのサンプルテキストです")
b = bedrock_embedding("味噌煮込みうどん。")
c = bedrock_embedding("This is Sample Text. for Experimental of embeddings.")

print(cosine_similarity(a, b))
print(cosine_similarity(b, c))
print(cosine_similarity(a, c))

実行前に、指定したリージョン（今回はus-west-2）のBedrockのコンソールからModel accessを開き、Titan Embeddings G1 - Textを有効にしておいてください。

有効化の方法は以下の記事を参考にしてください

実行結果

aとcが内容が近いので、3番目のコサイン類似度が最も高くなっています。日本語と英語でも比較できました。

0.3979828700786002
0.05232292622798426
0.5749125163290298

トークン数について

上記コードではコメントアウトされていますが、inputTextTokenCountにトークン数が入っています。 a=8, b=5, c=10という値になっていました。cは英語なので単語ごとに1トークンで計算は合っていますが、aとbの日本語でもトークン数が単語毎のように見えます。

日本語の場合、1文字1トークンになるモデルと比較するとお安く使えるかもしれません。

おわりに

Amazon BedrockのEmbeddingsを使ってみました。AWSで動かしているサービスと組み合わせて類似記事のレコメンデーションやベクトルデータベースと組み合わせてRAGなどの実装に有用に使えそうです。

Amazon Titanモデルはほかにもテキスト生成などバリエーションが増えていくようなので楽しみですね。