Amazon Titan Text Embedding V2とPythonでPostgreSQL+pgvectorのベクトルデータベースに゙保存してみた

生成AIの検索サービスにおいて、RAG(Retrieval-Augmented Generation)という仕組みを用いてユーザーのナレッジに基づいて回答を生成することがあります。

RAGでは、ナレッジとなるドキュメントをEmbedding(埋め込み)モデルでベクトル化してベクトルデータベースに保存し、ユーザーの問い合わせもEmbeddingでベクトル化して該当するドキュメントを検索し、検索したドキュメントをコンテキストに回答を生成します。非構造化データをベクトル表現という単一インターフェースで情報探索しています。

※ 図の引用元 https://aws.amazon.com/jp/blogs/news/knowledge-bases-for-amazon-bedrock-rag-patent/

本記事では、図の(1)~(4)のドキュメントをベクトル化して保存するステップについて、Pythonから埋め込みAPIを呼び出してPostgreSQL + 拡張モジュールpgvectorのベクトルデータベースに保存する方法を紹介します。

構成

環境

AWS上で環境構築します。

• AWSリージョン : US East (N. Virginia)
• ベクトルデータベース : Amazon RDS PostgreSQL 16.3 + pgvector 0.7.0
• テキスト埋め込みモデル : Amazon Titan Text Embedding V2
• ベクトルの次元 : 1024
• プログラミング言語 : Python 3.12(OSデフォルト)
• OS : EC2 Ubuntu 24.04

ウォークスルー

AWS リージョンの制約

今回利用する埋め込みモデルは 2024年6月4日時点ではAWSの一部のリージョンでしか利用できません。特に、東京リージョンでは利用できません。

以下のリージョンで各AWSリソースを作成しましょう。

• US East (N. Virginia)
• US West (Oregon)

ベクトルデータベースの作成

RDBにはAmazon RDS PostgreSQLを利用し、ベクトルデータベース化のために拡張モジュールpgvectorを利用します。

Auroraに比べてRDSはより新しいバージョンのDBエンジンや拡張モジュールがリリースされる傾向があります。 今回は2024年4月末にリリースされた pgvector 0.7.0 を使いたかったため、RDS PostgreSQLの最新版である16.3を採用しました。

pgvector の有効化やテーブル作成は後ほど行います。

EC2 の作成

ベクトルデータベース(PostgreSQL)の操作クライアントとなるEC2を用意します。

今回はAmazon EC2でUbuntu 24.04 を用意しました。

EC2に割り当てられているIAMロールにおいて、以下のアクションが許可(Allow)されていることを確認してください。

• bedrock:InvokeModel # モデルの呼び出し

EC2起動後は $ sudo apt update を実行します。

EC2からPostgreSQLへの接続

EC2 に PostgreSQL クライアントをインストールし、RDS PostgreSQL に接続します。

$ sudo apt install -y postgresql-client
$ psql -V
psql (PostgreSQL) 16.3 (Ubuntu 16.3-0ubuntu0.24.04.1)

$ HOST=rds-pg-16.xxx.us-east-1.rds.amazonaws.co
$ psql -h $HOST -U postgres test
Password for user postgres:
psql (16.3 (Ubuntu 16.3-0ubuntu0.24.04.1))
SSL connection (protocol: TLSv1.3, cipher: TLS_AES_256_GCM_SHA384, compression: off)
Type "help" for help.

test=> select version();
                                                   version
--------------------------------------------------------------------------------------------------------------
 PostgreSQL 16.3 on aarch64-unknown-linux-gnu, compiled by gcc (GCC) 7.3.1 20180712 (Red Hat 7.3.1-6), 64-bit
(1 row)
test=>

ベクトルデータを格納できるように pgvector Extentionを有効化します。

test=> CREATE EXTENSION vector;
CREATE EXTENSION
test=> \dx

                             List of installed extensions
  Name   | Version |   Schema   |                     Description
---------+---------+------------+------------------------------------------------------
 plpgsql | 1.0     | pg_catalog | PL/pgSQL procedural language
 vector  | 0.7.0   | public     | vector data type and ivfflat and hnsw access methods
(2 rows)

ドキュメント用のテーブルを作成します。

次元(1024)は埋め込みモデルに合わせて調整してください。

test=> CREATE TABLE documents (
    id bigserial PRIMARY KEY,
    content text,
    embedding vector(1024)
);


test=> CREATE INDEX ON documents USING hnsw (embedding vector_l2_ops);
CREATE INDEX

test=> \d+ documents
                                                           Table "public.documents"
  Column   |     Type     | Collation | Nullable |                Default                | Storage  | Compression | Stats target | Description
-----------+--------------+-----------+----------+---------------------------------------+----------+-------------+--------------+-------------
 id        | bigint       |           | not null | nextval('documents_id_seq'::regclass) | plain    |             |              |
 content   | text         |           |          |                                       | extended |             |              |
 embedding | vector(1024) |           |          |                                       | external |             |              |
Indexes:
    "documents_pkey" PRIMARY KEY, btree (id)
    "documents_embedding_idx" hnsw (embedding vector_l2_ops)
Access method: heap

pgvector はデフォルトでは最近傍検索を行います。近似最近傍検索に対応した IVFFlat と HNSW(0.5から)の2種類のインデックスにも対応しています。

また、pgvector 0.7.0 以降はベクトルを表現する型の選択肢も増えています。

型	最大次元	追加バージョン
vector	2000	-
halfvec	4000	0.7.0
bit	64000	0.7.0
sparsevec	1000(0の数)	0.7.0

Amazon Titan埋め込みモデルの有効化

テキストのベクトル化にあたり、今回はAmazonが提供するテキスト埋め込みモデルAmazon Titan Text EmbeddingsのV2を利用します。

このモデルは2024年4月末に利用可能となり、以下の特徴を持っています。

• 最大8192トークン
• ベクトルの次元 : 256/512/1024(デフォルト)
• 日本語を含めた100以上の言語に対応
• 正規化対応

ベクトルの次元(サイズの大きさ)と精度はトレードオフの関係にあります。1024次元を100%としたときに、512次元は約99%、256次元は約97%の精度を提供します。

Amazon Bedrockサイドメニューの Model accessから、利用するモデルが利用可能なことを確認してください。

Available to request となっている場合は、リクエスト申請が必要です。

AWS CLIからモデルの詳細を確認することも可能です。

$ MODEL_ID=amazon.titan-embed-text-v2:0
$ aws bedrock get-foundation-model --model-identifier $MODEL_ID
{
    "modelDetails": {
        "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0",
        "modelId": "amazon.titan-embed-text-v2:0",
        "modelName": "Titan Text Embeddings V2",
        "providerName": "Amazon",
        "inputModalities": [
            "TEXT"
        ],
        "outputModalities": [
            "EMBEDDING"
        ],
        "responseStreamingSupported": false,
        "customizationsSupported": [],
        "inferenceTypesSupported": [
            "ON_DEMAND"
        ],
        "modelLifecycle": {
            "status": "ACTIVE"
        }
    }
}

PythonからAmazon Titan埋め込みモデルを呼び出す

有効化したAmazon Titanモデルを Amazon Bedrock の InvokeModel APIから呼び出します。

PythonはOSにシステムインストールされているものを利用します。バージョンは3.12です。

$ python3 -V
Python 3.12.3

次に依存するライブラリ群をインストールします。

$ sudo apt install -y python3-boto3 python3-pip python3-psycopg

aptパッケージが提供されていないpgvectorはpipからインストールします。

$ sudo pip3 install pgvector --break-system-packages

以下がAWS Python SDKからのミニマルコードです。 サービス名が bedrock ではなく bedrock-runtime であることにご注意ください。

import boto3
import json

bedrock_runtime = boto3.client("bedrock-runtime")

inputText = """\
Your text string goes here
"""

body = json.dumps(
    {
        "inputText": inputText,
        "dimensions": 1024,
        "normalize": True,
    }
)


modelId = "amazon.titan-embed-text-v2:0"
accept = "application/json"
contentType = "application/json"


response = bedrock_runtime.invoke_model(
    body=body, modelId=modelId, accept=accept, contentType=contentType
)

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

print(len(embedding)) # 1024
print(embedding[:10]) # [-0.1070714, 0.06840672, 0.03816897, 0.025776448, 0.046100184, 0.028378878, 0.045852333, 0.0032530373, 0.04213458, 0.027139625]

正規化された1024次元のベクトルが返ってきていますね。

PostgreSQLに゙保存

変数 documents で定義したチャンク化後の文字列を Titanテキスト埋め込みモデル(Amazon Titan Text Embedding V2)でベクトル化するのが、以下のスクリプトです。bedrock:InvokeModel のアクションを利用しているため、呼び出し元は同権限が必要なことに注意してください。

• PostgreSQLの接続情報(DSN)
• 次元(DIMENSION)
• モデル(MODEL_ID)

等は適宜読み替えてください。

import json

import boto3

import psycopg
from pgvector.psycopg import register_vector

DIMENSION = 1024 # ベクトル次元数
DSN = "host=rds-pg-163.XXX.us-east-1.rds.amazonaws.com dbname=test user=postgres password=123"
MODEL_ID = "amazon.titan-embed-text-v2:0"

# ベクトル化するドキュメント
documents = [
  "福岡市の水族館「マリンワールド海の中道」には「リロ」という名前のラッコがいます",
  "アフリカに生息するダチョウは二足歩行の動物の中では最速",
  "アフリカに生息するパタスモンキーは霊長類の中では最速",
]

bedrock_runtime = boto3.client("bedrock-runtime")

def get_embedding(text):


    body = json.dumps(
        {
            "inputText": text,
            "dimensions": DIMENSION,
            "normalize": True,
        }
    )


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


    response = bedrock_runtime.invoke_model(
        body=body,
        modelId=MODEL_ID
    )

    response_body = json.loads(response["body"].read())
    return response_body.get("embedding")


with psycopg.connect(DSN, autocommit=True) as conn:
    register_vector(conn)
    with conn.cursor() as cur:
        for document in documents:
            embedding = get_embedding(document)
            cur.execute(
                "INSERT INTO documents (content, embedding) VALUES (%s, %s)", (document, embedding)
            )

埋め込みAPIを呼び出さずに、まとまった量のベクトル化されたダミーデータを追加したい場合、以下の様なスクリプトを利用します。

import random

import psycopg
from pgvector.psycopg import register_vector

NUM = 10000      # ダミーレコード数
DIMENSION = 1024 # ベクトル次元数

DSN = "host=rds-pg-163.XXX.us-east-1.rds.amazonaws.com dbname=test user=postgres password=123"

with psycopg.connect(DSN, autocommit=True) as conn:
    register_vector(conn)
    with conn.cursor() as cur:
        for _ in range(NUM):
            embedding = [random.uniform(-1, 1) for _ in range(DIMENSION)]
            cur.execute(
                "INSERT INTO documents (embedding) VALUES (%s)", (embedding,)
            )

最後にINSERTしたレコードを確認します。

test=> SELECT id, content FROM documents ORDER BY id ;
 id |                                    content
----+--------------------------------------------------------------------------------
  1 | 福岡市の水族館「マリンワールド海の中道」には「リロ」という名前のラッコがいます
  2 | アフリカに生息するダチョウは二足歩行の動物の中では最速
  3 | アフリカに生息するパタスモンキーは霊長類の中では最速
(3 rows)

ドキュメントの類似度をコサイン類似度から確認

ID:2のダチョウに似たドキュメントを、ベクトルデータに対してコサイン類似度で確認します。

コサイン距離には演算子 <=>が割り当てられているため、1-コサイン距離からコサイン類似度が求まります。 コサイン類似度は-1から1の間の値を取り、1に近づくほど似て、-1に近づくほど似なくなります。

SELECT
    id,
    1 - (embedding <=> (SELECT embedding FROM documents WHERE id = 2)) AS cosine_similarity,
    content
FROM documents;

 id |  cosine_similarity  |                                    content
----+---------------------+--------------------------------------------------------------------------------
  1 | 0.08991735637680565 | 福岡市の水族館「マリンワールド海の中道」には「リロ」という名前のラッコがいます
  2 |                   1 | アフリカに生息するダチョウは二足歩行の動物の中では最速
  3 | 0.41030489863782593 | アフリカに生息するパタスモンキーは霊長類の中では最速
(3 rows)

同じドキュメントは1で完全一致し、アフリカに生息し、最速という共通点のあるID=3のパタスモンキーのドキュメントも0.4とそれなりの類似度となり、期待通りです。

最後に

PostgreSQL+pgvectorをバックエンドとするベクトルデータベースにAmazon Titanの埋め込みモデルをAmazon Bedrockから呼び出してドキュメントをベクトル化する方法を紹介しました。

Knowledge bases for Amazon Bedrockのようなサービスを利用すると、ドキュメントのベクトル化ややベクトルデータベースへの格納をフルマネージドで行ってくれるため、簡単にRAGを構築できます。

本ケースではマネージドサービスやライブラリ/フレームワークを極力利用せずに、各処理に対応するAPIを直接呼び出しているため、ベクトル化の流れが理解しやすいのではないかと思います。

それでは。

参考

• Amazon Titan Text Embeddings models - Amazon Bedrock
• GitHub - pgvector/pgvector: Open-source vector similarity search for Postgres