AWS SDK for JavaScript v3 で Amazon Titan Text Embeddings を試す

はじめに

AWS SDK for JavaScript v3 を利用して Amazon Titan Text Embeddings を試してみた、という小ネタです。

現時点で Bedrock 関連ドキュメントのコード例は Python で書かれていることが多く、JavaScript SDK を利用する際には実装例を解釈して書き直すというひと手間が必要になります。また、SDK ドキュメントに載っているコード例がミスリーティング [1] なことがあり、困惑の種となっていました。

今回は Amazon Titan Text Embeddings v2 モデルを利用して、文字列を Embedding する例を紹介します。

コード例

下記に簡単なコード例を示します。

import {
  BedrockRuntimeClient,
  InvokeModelCommand,
} from "@aws-sdk/client-bedrock-runtime";

// client を初期化
const bedrockClient = new BedrockRuntimeClient({ region: "us-east-1" });

// Embedding 対象の文字列を指定する
const inputText = "ここにベクトル化したい文字列を指定します。"

// Command の定義
const command = new InvokeModelCommand({
  modelId: "amazon.titan-embed-text-v2:0",
  contentType: "application/json",
  body: JSON.stringify({ inputText }),
});
const res = await bedrockClient.send(command);

// Uint8Array として返ってくる Response body を json object にパースする
const decoder = new TextDecoder("utf-8");
const jsonObject = JSON.parse(decoder.decode(res.body)) as {
  embedding: number[];
  inputTextTokenCount: number;
};

ポイントとしては以下のとおりです:

• InvokeModelCommand を実行する際は contentType に application/json を指定し、 body として { inputText: string } を stringify したものを渡す
• API の Response Body は Uint8Array として返ってくるため TextDecoder を利用してデコード、パースする

上記のコードを実行することで、対象のテキストを Embedding した結果のベクトル（数値の配列）を取得することができます。このベクトルを、ベクトル検索 (Vector Search) の Index に指定することで、ベクトルデータベースを構築することができます。あとは、RAG のソースにするなど、使い方は様々ですね。

以上、AWS SDK for JavaScript v3 を利用して Amazon Titan Text Embeddings モデルを試してみた、というお話でした。