PostgreSQLにはベクトルデータを扱えるようにするために、pgvectorという拡張機能があります。
これによって、既存のリレーショナル形式のデータとベクトルデータを同じ表内で利用できます。
ベクトルデータのみ専用のDBを設けて、そちらに格納する方法も考えられますが、
開発効率やDB運用の手間を考えると、同じDB上に集約して使えるほうが良さそうです。
ということで、今回はPostgreSQLを使い、既存のリレーショナル表に対して
pgvectorを使ってベクトルデータを格納する方法を検証してみました。
事前準備
PostgreSQL、およびPython環境の事前準備については以下の記事を参考にしてください。
生成AI (Cohere)＋LangChain＋Vector Database (PostgreSQL)でRAGを試してみた
上記の「必要なPythonパッケージのインストール」から「PostgreSQLへのホスト名指定でのログインを有効化」までが対象作業です。
環境が整ったら、検証に使う表を作成します。
CREATE TABLE products_v (
product_id SERIAL PRIMARY KEY,
product_name VARCHAR(255),
manufacturer_name VARCHAR(255),
product_overview TEXT,
price FLOAT,
release_date DATE,
embedding vector(4096)
);
サンプルデータを入れます。
INSERT INTO products_v (product_name, manufacturer_name, product_overview, price, release_date)
VALUES ('Tiny Tunes Piano',
'Playful Minds Co.',
'A colorful and musical toy piano that lets toddlers play and learn songs. The piano has 12 keys, each with a different animal sound and a corresponding sticker. The piano also has a mode switch that allows toddlers to choose between animal sounds, musical notes, or songs. The piano comes with a songbook that teaches toddlers how to play 10 popular nursery rhymes.',
29.99,
'2023-12-01');
INSERT INTO products_v (product_name, manufacturer_name, product_overview, price, release_date)
VALUES ('Tiny Blocks Castle',
'Playful Minds Co.',
'A set of 50 wooden blocks in various shapes and colors that toddlers can use to build their own castle. The blocks are easy to stack and connect with magnets. The set also includes a drawbridge, a flag, and two knight figures. The blocks are made of natural and eco-friendly materials and are safe for toddlers to play with.',
39.99,
'2024-01-15');
INSERT INTO products_v (product_name, manufacturer_name, product_overview, price, release_date)
VALUES ('Tiny Pals Farm',
'Playful Minds Co.',
'A fun and interactive toy farm that introduces toddlers to farm animals and their sounds. The farm has a barn, a silo, a tractor, and a fence. The farm also has 10 animal figures, each with a button that plays its sound when pressed. The farm has a sensor that recognizes the animal figures and plays a song or a fact about them when placed on the farm.',
49.99,
'2024-02-01');
Pythonスクリプトの作成
既存の products_v 表の product_overview 列の値を使い、Cohereのモデルを使ってベクトルデータを生成します。
出来たベクトルデータをUPDATE文で products_v 表の embedding 列に格納します。
ちなみに embedding 列にベクトルデータを格納する際は、CohereのAPIから返ってきた配列のベクトルデータを、NumPy配列に変換して格納する必要があります。
import os
import numpy as np
import psycopg2
from pgvector.psycopg2 import register_vector
from langchain.vectorstores.pgvector import PGVector
from langchain.embeddings import CohereEmbeddings
COHERE_API_KEY = os.environ['COHERE_API_KEY']
embeddings = CohereEmbeddings(cohere_api_key=COHERE_API_KEY)
# データベースへの接続情報を設定します
conn = psycopg2.connect(
host="localhost",
database="postgres",
user="postgres",
password="postgres"
)
register_vector(conn)
# カーソルを作成します
cur = conn.cursor()
# SQLクエリを実行します
cur.execute("SELECT product_id, product_overview FROM products_v;")
# 結果を取得します
rows = cur.fetchall()
# 各行のproduct_overviewをベクトル化してproducts_vに格納します
for row in rows:
doc_result = embeddings.embed_documents([row[1]])
cur.execute('UPDATE products_v SET embedding = %s WHERE product_id = %s;', (np.array(doc_result[0]), row[0]))
# カーソルと接続を閉じます
conn.commit()
cur.close()
conn.close()
上記コードを embedding.py として保存します。
実行結果
コードを実行します。
psqlで products_v 表の embedding 列を確認します。
SELECT product_id, embedding FROM products_v;
product_id | embedding
------------+----------------------------------------------------------------------------------------
1 | [-0.31567383,-2.484375,0.58935547,2.5546875,-1.5625,0.2590332,2.1933594,1.703125,0.8295 ...
2 | [1.0498047,-0.9082031,-1.8886719,0.9897461,-2.6992188,-1.2197266,1.921875,1.0507812,-0. ...
3 | [1.6728516,-0.46484375,-0.9838867,2.4414062,0.55029297,-0.9790039,0.78808594,1.7128906, ...
以上のようにすれば、リレーショナルデータとベクトルデータを同じ表で扱えます。
データモデルに応じて種類の異なるDBを構築・運用する必要がなくなり、開発効率や運用効率の向上に繋がるのではないでしょうか。
Slowly But Surely 