【AWS】Bedrock で Titan から文書埋め込みを取得する

AWS の Bedrock と、基盤モデルの Titan を使って文書埋め込みの取得をします。

コード (Python3)

import json
import boto3
import numpy as np


class Bedrock:
    def __init__(
        self,
    ):
        self.model_id = "amazon.titan-embed-text-v1"

    def completion(self, text):
        session = boto3.Session(region_name="ap-northeast-1")

        bedrock = session.client(service_name="bedrock-runtime")

        body = json.dumps({"inputText": text})

        accept = "application/json"
        contentType = "application/json"

        response = bedrock.invoke_model(
            body=body, modelId=self.model_id, accept=accept, contentType=contentType
        )

        response_body = json.loads(response.get("body").read())

        return response_body["embedding"]


def test():
    bedrock = Bedrock()
    response_1 = bedrock.completion("Today is Friday in California.")
    response_2 = bedrock.completion("In California, it is Friday today.")
    response_3 = bedrock.completion("I am a pen.")

    # calculate these euclid distance
    return {
        "1 vs 2": np.linalg.norm(np.array(response_1) - np.array(response_2)),
        "1 vs 3": np.linalg.norm(np.array(response_1) - np.array(response_3)),
        "2 vs 3": np.linalg.norm(np.array(response_2) - np.array(response_3)),
    }


print(json.dumps(test(), indent=2))

test() では、以下の 3 つの文章のユークリッド距離を計算しています。

1. Today is Friday in California.
2. In California, it is Friday today.
3. I am a pen.

意味的には 1 と 2 の距離は近くて、3 は他の 2 つと遠い気がします。

実行結果

{
  "1 vs 2": 6.021247346519345,
  "1 vs 3": 22.76307302608681,
  "2 vs 3": 22.32505996689041
}

期待通りの結果となりました。

応用

Titan Embedding の埋め込み結果がちゃんと正規化されていればユークリッド距離を文書類似度とすることもできると思われます。
上記の workshop のように、埋め込みを FAISS に突っ込んで RAG にするという応用があります。

おわり

Titan Embedding は補完用のふつうの Titan とは違ってマルチリンガルらしいです。日本語も使える～やった～（ほんとか？）