はじめに
この記事はNTTテクノクロス Advent Calendar 2023 シリーズ1の25日目（最終日）です。
NTTテクノクロスの金澤です。
普段はPostgreSQL関連の開発や検証業務に携わっています。
今回はPostgreSQLの拡張機能pgvectorとOpen AIのAPIを利用してベクトルデータを扱うデモ環境を作成して動かしてみたいと思います。
pgvectorとは何か
pgvectorとは、PostgreSQL上でベクトルデータ型(vector)を保存し、検索する機能を追加する拡張機能になります。
この拡張機能を使用することで、RDBMSであるPostgreSQLにベクトルデータベースの機能を追加することができます。
先日行われたPostgreSQL Conference Japan 2023ではpgvectorの講演が行われていたり、AWS/Azure/GCPなどの各メガクラウドが提供しているDBサービスでもpgvector拡張機能に対応していたりと、近年注目されている拡張機能の一つであります。
構築するデモ環境について
今回は、Dockerコンテナでpgvector拡張機能をもつPostgreSQLデータベースを作成します。
vector型のカラムを持つテーブルに、OpenAI が提供しているAPIでベクトルデータ型に変換(Embedding)した文章を保存し、検索操作を実行していきます。
pgvector拡張機能をもつPostgreSQLとOpen AI APIを組み合わせた環境の構築については、以下のサイトを参考にしました。
また、今回使用する文章のデータセットやデータのセットアップ・操作等は、以下Google Cloud blogの内容を使用いたしました。
上記のGoogle Cloud blogで使用しているデータセットを利用し、約 800 のおもちゃ関連商品のデータを保存するテーブルを作成します。
環境を構築する
上述のサイトを参考に、今回はDockerコンテナで構築したPostgreSQLに、Jupyter Notebooksから接続していきます。
pgvector拡張機能をもつPostgreSQLをDocker コンテナで起動する
pgvector 開発者の方が提供しているDockerイメージを使用して、pgvector拡張機能をもつPostgreSQLを起動します。
$ docker pull ankane/pgvector
$ docker run --name pgvector-demo -e POSTGRES_PASSWORD=pswd -p 5432:5432 -d ankane/pgvector
起動したPostgreSQLのコンテナは、ホストのポート5432にマッピングされているため、以下のようにpsqlを実行することでDBへ接続することができます。
$ psql -p 5432 -h localhost -U postgres
DBへ接続することができたら、｀CREATE DATABASEでvector_dbを作成し、CREATE EXTENSIONでpgvector拡張機能を有効にします。
postgres=# CREATE DATABASE vector_db;
CREATE DATABASE
vector_db=# CREATE EXTENSION vector;
CREATE EXTENSION
pgvector拡張機能を有効にしたら、2つのテーブルを作成します。
商品名や値段を保存するproductsテーブルと商品の説明文と説明文をベクトルに変換したvenctor型のカラムをもつproduct_embeddingsテーブルです。
Open AI APIでは、ベクトルの次元は1536となるので、vector(1536)と定義します。
CREATE TABLE products(
product_id VARCHAR(1024) PRIMARY KEY,
product_name TEXT,
description TEXT,
list_price NUMERIC);
CREATE TABLE product_embeddings(
product_id VARCHAR(1024) NOT NULL REFERENCES products(product_id),
content TEXT,
embedding vector(1536));
以上でデータベース側の設定は完了になります。
テーブルにデータを入力する
Jupyter Notebooksで作成したPostgreSQLに接続して、テーブルにデータを入力していきます。
事前に、OpenAI APIのAPI Keyを取得し、.envファイルに保存しておきます。
OPENAI_API_KEY=....
必要なライブラリをインストール、インポートします。
!pip install pandas asyncio asyncpg langchain openai python-dotenv
import pandas as pd
import asyncio
import asyncpg
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from dotenv import load_dotenv
load_dotenv()
続いて、Google Cloud blogで使用していたサンプルのデータセットを保存します。
# データセットをダウンロードして Pandas データフレームに保存する。
DATASET_URL='https://github.com/GoogleCloudPlatform/python-docs-samples/raw/main/cloud-sql/postgres/pgvector/data/retail_toy_dataset.csv'
df = pd.read_csv(DATASET_URL)
df = df.loc[:, ['product_id', 'product_name', 'description', 'list_price']]
df.head(5)
データセットには以下のようなデータが含まれています。
| product_id | product_name | description | list_price | |
|---|---|---|---|---|
| 0 | 7e8697b5b7cdb5a40daf54caf1435cd5 | Koplow Games Set of 2 D12 12-Sided Rock, Paper... | Rock, paper, scissors is a great way to resolv... | 3.56 |
| 1 | 7de8b315b3cb91f3680eb5b88a20dcee | 12"-20" Schwinn Training Wheels | Turn any small bicycle into an instrument for ... | 28.17 |
| 2 | fb9535c103d7d717f0414b2b111cfaaa | Bicycle Pinochle Jumbo Index Playing Cards - 1... | Purchase includes 1 blue deck and 1 red deck. ... | 6.49 |
| 3 | c73ea622b3be6a3ffa3b0b5490e4929e | Step2 Woodland Adventure Playhouse & Slide | The Step2 Woodland Climber Adventure Playhouse... | 499.99 |
| 4 | dec7bd1f983887650715c6fafaa5b593 | Step2 Naturally Playful Welcome Home Playhouse... | Children can play and explore in the Step2 Nat... | 600.00 |
データセットの情報をproductsテーブルにcopyします。
conn = await asyncpg.connect(
user='postgres',
password='pswd',
database='vector_db',
host='localhost',
port=5432
)
tuples = [tuple(str(value) for value in row) for row in df.itertuples(index=False)]
# ローカルの PostgreSQL にデータをコピーする
await conn.copy_records_to_table('products', records=tuples, columns=list(df), timeout=10)
一度、psqlコマンドでPostgreSQLに接続し、productsテーブルのデータ数を確認します。
vector_db=# SELECT count(*) FROM products;
count
-------
850
(1 row)
850件のデータが保存されていることを確認できました。
続いては、description列に含まれる商品説明文を分割して、ベクトルデータに変換します。
from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(
separators=[".", "\n"],
chunk_size=500,
chunk_overlap=0,
length_function=len,
)
chunked = []
for index, row in df.iterrows():
product_id = row['product_id']
desc = str(row['description'])
splits = text_splitter.create_documents([desc])
for s in splits:
r = {'product_id': product_id, 'content': s.page_content}
chunked.append(r)
embeddings_service = OpenAIEmbeddings()
batch_size = 5
for i in range(0, len(chunked), batch_size):
request = [x['content'] for x in chunked[i: i + batch_size]]
response = embeddings.embed_documents(request)
for x, e in zip(chunked[i : i + batch_size], response):
x["embedding"] = e
product_embeddings = pd.DataFrame(chunked)
product_embeddings.head()
product_embeddings.head()で出力した結果は以下となり、embeding列にベクトルデータに変換された値が格納されています。
| product_id | content | embedding | |
|---|---|---|---|
| 0 | 7e8697b5b7cdb5a40daf54caf1435cd5 | Rock, paper, scissors is a great way to resolv... | [-0.00886749972743081, 0.014060657961181389, 0... |
| 1 | 7e8697b5b7cdb5a40daf54caf1435cd5 | . Great for educational games, dice games, boa... | [-0.0013363465225081973, 0.02269774532722539, ... |
| 2 | 7de8b315b3cb91f3680eb5b88a20dcee | Turn any small bicycle into an instrument for ... | [0.01945479213474804, -0.00412737445352915, 0.... |
| 3 | 7de8b315b3cb91f3680eb5b88a20dcee | . Durable Construction: Steel brackets stand u... | [0.014991591231696421, 0.00926185557582806, -0... |
| 4 | 7de8b315b3cb91f3680eb5b88a20dcee | . Tools required: Adjustable wrench. www.schwi... | [0.02521163542924525, -0.001772273162437382, -... |
変換したベクトルデータはproduct_embeddingsにINSERTしていきます。
from pgvector.asyncpg import register_vector
import numpy as np
await register_vector(conn)
for index, row in product_embeddings.iterrows():
await conn.execute(
"INSERT INTO product_embeddings (product_id, content, embedding) VALUES ($1, $2, $3)",
row["product_id"],
row["content"],
np.array(row["embedding"]),
)
実行後、PostgreSQLに接続して、product_embeddingsのデータ数をカウントして、データが追加されていることを確認します。
vector_db=# SELECT count(*) FROM product_embeddings ;
count
-------
2671
(1 row)
以上で、テキストデータをベクトルに変換して、PostgreSQL上にあるテーブルへの保存が完了し環境構築が終了しました。
データを検索してみる
続いては、データベースにクエリを投げて、動作を確認してみます。
pgvectorでは、L2距離(<->),内積(<#>),コサイン距離(<=>)によるベクトル距離比較がサポートされています。
それぞれの距離比較の特徴については、前述したPostgreSQL Conference Japan 2023 発表資料で大変わかりやすくまとめらています。
関心のある方は、当該資料をご確認ください。
今回は、Google Cloud blog内でも使用していたコサイン距離(<=>)によるベクトル距離比較を実行してみます。
<=>でベクトル間のコサイン距離を比較することで、ベクトルの向きの類似度を算出できます。
コサイン類似度はドュメント検索でよく利用されます。
検索する条件として、"Do you have a beach toy set that teaches numbers and letters to kids?"というワードをベクトルに変換し、コサイン類似度を比較します。
user_query = "Do you have a beach toy set that teaches numbers and letters to kids?"
vector = embeddings.embed_query(user_query)
print (vector)
print (vector)の結果を確認すると、以下のようなベクトルデータに変換されています。
[-0.008257456479335568, -0.01240224910278451, 0.008906485896301887, -0.022902392836960633, -0.029688286941888034, 0.04032980079236429, -0.011476899991323739, -0.0008803669276621656, -0.00988966921081363, 0.0012844038981710644, 0.0059183791965176005, 0.0216300366445488, ...]
コサイン類似度を使ったベクトル検索
以下のようなSQLクエリを作成して、検索を実行します。
--$1 ベクトルデータ
--$2 検索時の閾値条件
--$3 取得するデータ数
WITH vector_matches AS (
SELECT
product_id,
MAX(1 - (embedding <=> $1)) AS similarity
FROM product_embeddings
WHERE 1 - (embedding <=> $1) > $2
GROUP BY product_id
ORDER BY similarity DESC
LIMIT $3
)
SELECT
products.product_id,
products.product_name,
products.list_price,
products.description,
vector_matches.similarity
FROM
products
JOIN
vector_matches ON products.product_id = vector_matches.product_id
ORDER BY
vector_matches.similarity DESC;
質問文のベクトルデータはvector変数内にあるため、Jupyter Notebook側からSQLを実行します。
matches = []
similarity_threshold = 0.1
num_matches = 50
results = await conn.fetch(
"""
WITH vector_matches AS (
SELECT
product_id,
MAX(1 - (embedding <=> $1)) AS similarity
FROM product_embeddings
WHERE 1 - (embedding <=> $1) > $2
GROUP BY product_id
ORDER BY similarity DESC
LIMIT $3
)
SELECT
products.product_id,
products.product_name,
products.list_price,
products.description,
vector_matches.similarity
FROM
products
JOIN
vector_matches ON products.product_id = vector_matches.product_id
ORDER BY
vector_matches.similarity DESC;
""",
vector,
similarity_threshold,
num_matches,
)
if len(results) == 0:
raise Exception("Did not find any results. Adjust the query parameters.")
for r in results:
matches.append(
{
"product_id": r["product_id"],
"product_name": r["product_name"],
"description": r["description"],
"list_price": round(r["list_price"],2),
"similarity": r["similarity"]
}
)
matches = pd.DataFrame(matches)
matches.head(5)
上位5つには以下のような結果を得ることができました。
(similarityの値が1に近ければ近いほど似ていると判断する)
| product_id | product_name | description | list_price | similarity | |
|---|---|---|---|---|---|
| 0 | 20e597d8836d9e5fa29f2bd877fc3e0a | CoolSand Learning Sand Molds and Tools Kit (44 Pcs) - Wo... | Start Your Kids Creativity Now! Treat your little ones ... | 31.00 | 0.874993 |
| 1 | a57aced63b89c8b5f0a53280c10e47d1 | Beach Toys Mega Pack - Bulk Beach & Sandbox Play Set w/ ... | Keep them busy at the beach?while molding their minds! A... | 10.99 | 0.872064 |
| 2 | 5a91b29911ab59c8de7668df262e7231 | Bisontec Beach Toy Playset With Wheelbarrow (Colors May ... | Toy Playset Beach Toy Great Gift for Ages 3 and up. This... | 14.48 | 0.860930 |
| 3 | e87c6524f0b88aab70f11d0234f74cdf | Kids Beach Sand Toys Set for Gift with Sand Molds,Mesh B... | Educational: Because sand is an excellent tool for enha... | 22.99 | 0.860120 |
| 4 | 445108acb8d2304938de7835b4f8b1fd | Beach Toys Bucket And Shovel Set - 8 Piece In Zipper Bag... | Are you looking for a perfect baby toy play set for beac... | 19.99 | 0.859285 |
WHERE句でフィルタリングした検索
上記で実行したクエリについて、WHERE句で条件を指定してすることも可能です。
今回は、productsテーブルにある料金データlist_priceで20ドル以上100ドル以下という条件で検索してみます。
--$4が最低金額(20ドル)/$5が最高金額(100ドル)
WITH vector_matches AS (
SELECT
product_id,
MAX(1 - (embedding <=> $1)) AS similarity
FROM product_embeddings
WHERE 1 - (embedding <=> $1) > $2
GROUP BY product_id
ORDER BY similarity DESC
LIMIT $3
)
SELECT
products.product_id,
products.product_name,
products.list_price,
products.description,
vector_matches.similarity
FROM
products
JOIN
vector_matches ON products.product_id = vector_matches.product_id
WHERE list_price >= $4 AND list_price <= $5
ORDER BY
vector_matches.similarity DESC;
results内のクエリを変更して実施。
matches = []
results = await conn.fetch(
"""
WITH vector_matches AS (
SELECT
product_id,
MAX(1 - (embedding <=> $1)) AS similarity
FROM product_embeddings
WHERE 1 - (embedding <=> $1) > $2
GROUP BY product_id
ORDER BY similarity DESC
LIMIT $3
)
SELECT
products.product_id,
products.product_name,
products.list_price,
products.description,
vector_matches.similarity
FROM
products
JOIN
vector_matches ON products.product_id = vector_matches.product_id
WHERE list_price >= $4 AND list_price <= $5
ORDER BY
vector_matches.similarity DESC;
""",
vector,
similarity_threshold,
num_matches,
min_price,
max_price,
)
if len(results) == 0:
raise Exception("Did not find any results. Adjust the query parameters.")
for r in results:
matches.append(
{
"product_id": r["product_id"],
"product_name": r["product_name"],
"description": r["description"],
"list_price": round(r["list_price"],2),
"similarity": r["similarity"],
}
)
matches = pd.DataFrame(matches)
matches.head(5)
list_price"の金額が、指定した範囲のものになっていることが確認できます。
| product_id | product_name | description | list_price | similarity | |
|---|---|---|---|---|---|
| 0 | 20e597d8836d9e5fa29f2bd877fc3e0a | CoolSand Learning Sand Molds and Tools Kit (44 Pcs) - Wo... | Start Your Kids Creativity Now! Treat your little ones ... | 31.00 | 0.874993 |
| 1 | e87c6524f0b88aab70f11d0234f74cdf | Kids Beach Sand Toys Set for Gift with Sand Molds,Mesh B... | Educational: Because sand is an excellent tool for enha... | 22.99 | 0.860120 |
| 2 | c29104ac1a4e0faeabae7c8a32e9957b | American Plastic Toys Sand and Water Play Table | American Plastic Toys Sand and Water Play Table will pro... | 24.99 | 0.852799 |
| 3 | 1d15aa5cceefa670f0cc83de517dc4f9 | Melissa & Doug Jumbo ABC-123 Rug (58 x 79 inches – 36 Ga... | Kids will have jumbo amounts of fun exploring and learni... | 54.99 | 0.851722 |
| 4 | 2eeba67dba7c262da27f05622fc4b705 | Soft Neon 18 Piece Playset | 10 months & up. An ideal toy for every stage of developm... | 59.95 | 0.831987 |
全文検索(キーワード検索)と組み合わせて検索
PostgreSQLには、全文検索機能が備わっています。
キーワードを問い合わせて、文書内に該当する文章が存在するかを検索できます。
PostgreSQLに組み込まれている全文検索用の関数のうち、ユーザーが一般的なウェブ検索クエリのスタイルで入力することができるwebsearch_to_tsquery関数を今回は使用します。
キーワード検索する内容として、3 yearsというテキストを含むもの検索します。
PostgreSQLの全文検索クエリ例
-- $1 は'3 years'
SELECT
product_id,
product_name,
list_price,
description,
ts_rank_cd(to_tsvector('english', description), websearch_to_tsquery('english', $1)) AS rank
FROM
products
WHERE
to_tsvector('english', description) @@ websearch_to_tsquery('english', '$1')
ORDER BY
試しに、ベクトル検索は使用せずに、全文検索のみを実行してみます。
mmatches = []
keyword = '3 years'
full_text_results = await conn.fetch(
"""
SELECT
product_id,
product_name,
list_price,
description,
ts_rank_cd(to_tsvector('english', description), websearch_to_tsquery('english', $1)) AS rank
FROM
products
WHERE
to_tsvector('english', description) @@ websearch_to_tsquery('english', '$1')
ORDER BY
rank DESC
""",
keyword
)
# 結果の表示
for r in full_text_results:
matches.append(
{
"product_id": r["product_id"],
"product_name": r["product_name"],
"description": r["description"],
"list_price": round(r["list_price"],2),
"rank": r["rank"],
}
)
print("全文検索による絞り込み:")
print(len(matches))
matches = pd.DataFrame(matches)
matches.head()
print(len(matches))で絞り込んだデータ数を表示しており、313件となっています。
全文検索による絞り込み:
313
絞り込まれたデータを出力してみたところ以下のような結果になりました。
| product_id | product_name | description | list_price | rank | |
|---|---|---|---|---|---|
| 0 | 501257f1b3df0af591d939f53692cf77 | Ktaxon Tree Swing Disc Rope Swing 6.5ft - with Leg Safet... | FUNNY USE. 6.5ft high performance polyethylene rope wit... | 25.59 | 0.415703 |
| 1 | 45de24a6c20856ed2627b66995da7318 | OTVIAP Swimming Pool Toy,4pcs Swimming Pool Toys Mine Sh... | Description: This swimming toy will bring lots of fun i... | 10.68 | 0.304350 |
| 2 | 7e52aee88a6a6a0764cb68a69b2d1dc0 | Mini Lucky Jackpot Slot Machine Fun Gift Fruit Slot Mach... | Mini Lucky Jackpot Slot Machine Fun Gift Fruit Slot Mach... | 10.55 | 0.205263 |
| 3 | 8195e328afe639ae85e728b380e38e5b | Kids Indoor Inflatable Large Castle Sand Box Sandbox Sli... | 60*45cm Kids IndoorInflatable Large Castle Sand Box San... | 9.24 | 0.200885 |
| 4 | f933a3aab20a8739aa70e9d1f9ba278f | 7Pcs Two Colors Polyhedral Dice with Free Pouches for Du... | Specifications: Material: Acrylic Side Size: 16mm to 20... | 7.23 | 0.200813 |
上記の絞り込んだ結果に対して、ベクトル検索をすることで、キーワード検索と組み合わせてみます。
以下のように、全文検索を実施した結果をJOIN句で結合して、ベクトル検索を実施します。
SELECT
p.product_id,
MAX(p.rank) AS rank,
MAX(1 - (pe.embedding <=> $1)) AS similarity
FROM
product_embeddings pe
JOIN
(
SELECT
product_id,
ts_rank_cd(to_tsvector('english', description), websearch_to_tsquery('english', $4)) AS rank
FROM
products
WHERE
to_tsvector('english', description) @@ websearch_to_tsquery('english', $4)
) AS p ON pe.product_id = p.product_id
WHERE
1 - (pe.embedding <=> $1) > $2
GROUP BY
p.product_id
ORDER BY
similarity DESC
LIMIT $3
実行する内容
matches = []
keyword = '3 years'
results = await conn.fetch(
"""
WITH vector_matches AS (
SELECT
p.product_id,
MAX(p.rank) AS rank,
MAX(1 - (pe.embedding <=> $1)) AS similarity
FROM
product_embeddings pe
JOIN
(
SELECT
product_id,
ts_rank_cd(to_tsvector('english', description), websearch_to_tsquery('english', $4)) AS rank
FROM
products
WHERE
to_tsvector('english', description) @@ websearch_to_tsquery('english', $4)
) AS p ON pe.product_id = p.product_id
WHERE
1 - (pe.embedding <=> $1) > $2
GROUP BY
p.product_id
ORDER BY
similarity DESC
LIMIT $3
)
SELECT
products.product_id,
products.product_name,
products.list_price,
products.description,
vector_matches.similarity,
vector_matches.rank
FROM
products
JOIN
vector_matches ON products.product_id = vector_matches.product_id
ORDER BY
similarity DESC, rank DESC
LIMIT $3;
""",
vector,
similarity_threshold,
num_matches,
keyword
)
if len(results) == 0:
raise Exception("Did not find any results. Adjust the query parameters.")
for r in results:
matches.append(
{
"product_id": r["product_id"],
"product_name": r["product_name"],
"description": r["description"],
"list_price": round(r["list_price"],2),
"similarity": r["similarity"],
"rank": r["rank"]
}
)
matches = pd.DataFrame(matches)
matches.head()
検索結果を見ると、2以下の値が先ほど取得したベクトル検索結果と異なっていることが確認できます。
| product_id | product_name | description | list_price | similarity | rank | |
|---|---|---|---|---|---|---|
| 0 | 20e597d8836d9e5fa29f2bd877fc3e0a | CoolSand Learning Sand Molds and Tools Kit (44 Pcs) - Wo... | Start Your Kids Creativity Now! Treat your little ones ... | 31.00 | 0.874993 | 0.100000 |
| 1 | e39855540f05c9b350c6b04e6d7b1830 | 6Pcs Sand Kids Beach Toys Castle Bucket Spade Shovel Rak... | 6Pcs Sand Kids Beach Toys Castle Bucket Spade Shovel Ra... | 10.98 | 0.846966 | 0.100000 |
| 2 | 4fb5495d3825cd87fc1475af962a6eb7 | 11pc Kids Beach Toys Set Molds Tools, Sandbox Toys On Su... | 11pc KidsBeach Toys Set Molds Tools, Sandbox Toys On Su... | 12.55 | 0.840797 | 0.009221 |
| 3 | 2eeba67dba7c262da27f05622fc4b705 | Soft Neon 18 Piece Playset | 10 months & up. An ideal toy for every stage of developm... | 59.95 | 0.831987 | 0.200000 |
| 4 | 0e08271603b53b963f81dbdd02f989a9 | Spiralin Seas Waterpark | The Spiralin Seas Waterpark provides toddlers a splashin... | 41.98 | 0.829780 | 0.201538 |
まとめ
pgvectorでPostgreSQLにベクトルデータを入れることで、
ベクトルデータによる検索と既存のデータベースに登録している情報を使用したフィルタリングや、全文検索と組み合わせた検索を行うことができました。
条件の絞り込みやテーブルのJOINなど、RDBMSの機能も合わせて使用できるのがpgvectorならではの特徴だと思います。
今回はデモ環境に取り込んだデータ数が850程度と少なく、とりあえず動かしみるというところを目的としていたため、インデックスの設定などを行っておりません。
pgvectorはベクトル用のインデックスであるIVFFlat,HNSWが用意されており、うまく設定することで検索性能に大きく影響します。
今後は、大規模なデータセットでインデックスを追加等の検証も行っていきたいと考えています。
Advent Calendar2023に最後までお付き合いいただき、ありがとうございました！
参考
pgvector
pgvectorを使ってChatGPTとPostgreSQLを連携してみよう！（PostgreSQL Conference Japan 2023 発表資料）
VECTOR DATABASES - PGVECTOR AND LANGCHAIN
pgvector、LLM、LangChain を使用して Google Cloud データベースで AI 搭載アプリを構築する