pgvector入門 ~データ編~
はじめに
こんにちは。ミスミグループ本社 Gateway推進本部の飯塚です。レコメンドや検索、DB構築などを幅広く担当しています。
本記事は【pgvector 入門】と題して、PostgreSQL の環境構築からベクトルデータの格納、アプリ実装までを三部構成で連載していきます。
今回は ~データ編~ になります。LangChain を使ってドキュメントのベクトル化とデータベースへの格納を行い、pgvector が提供している近似計算について紹介します。
振り返り
前回は Amazon Aurora PostgreSQL Compatible 環境を構築し、Amazon SageMaker Code Editor を使って PostgreSQL へアクセスしました。
今回これらの環境を用いてドキュメントのベクトル化と格納を行いますので、環境構築が必要な方は前回記事をご参照ください。
前提知識
ドキュメントのベクトル化に際し、LangChain と Hugging Face を利用します。
これらは最近の機械学習、AIの文脈で必ずと言ってもいいほど出てくるフレームワークになりますので、しっかり押さえておきましょう。
LangChain とは
LangChain is a framework for developing applications powered by large language models (LLMs).
出典: https://python.langchain.com/docs/introduction/
LangChain とは、大規模言語モデル(LLM)を利用したアプリケーションの開発に使われれるフレームワークになります。Python および TypeScript に対応しており本記事では Python を用います。
LangChain を使う最大のメリットは、利用する生成AIモデルを自由に選択出来る点にあります。またインターフェースが共通化されていることにより、利用するモデルに依存することなくコーディングが可能です。加えて、生成AIアプリケーションによく使われる手法である Embedding や ReAct、RAG などが実装されているため、簡単にアプリケーション開発まで行えます。
今回は LangChain をインターフェースとして使い、ドキュメントのベクトル化(Embedding)と pgvector 環境へのデータ格納を行ってきます。
Hugging Face とは
The platform where the machine learning community collaborates on models, datasets, and applications.
出典: https://huggingface.co/
Hugging Face とは、機械学習に関するモデルやデータセット、アプリケーションが共有・公開されているプラットフォームです。モデルを一つ例にとっても 2025年2月現在 1,421,048 もの様々なモデルが公開されています。今回はこの中から embedding 用にモデルを使用します。
そしてもちろん LangChain は Hugging Face 用のインターフェースも用意していますので、たったの数行で Hugging Face にアップロードされているモデルを使用することが出来ます。
本題
では、上記 LangChain と Hugging Face を使ってベクトルデータを作成していきましょう。
大きく以下の流れで進めていきます。
- サンプルデータの準備
- python 環境の構築
- データのベクトル化および格納
- データ確認
- 近似計算
まずはサンプルデータの準備です。
サンプルデータの準備
今回は Amazon Web Services の概要 - AWS ホワイトペーパーを使っていきます。
リンク先からPDFファイルをダウンロードしておいてください。あとで Code Editor インスタンスにアップロードします。
python 環境の構築
前回作成した SageMaker Code Editor を引き続き使っていきます。
まずは Code Editor インスタンス上で必要なライブラリをインストールします。
pip install langchain_community
pip install langchain_huggingface
pip install langchain_postgres
pip install langchain-text-splitters
pip install pypdf
※パッケージの依存関係でエラーになった時は -U オプションを付けると回避出来る事があります
データのベクトル化および格納
データのベクトル化および格納は LangChain が提供している関数 PGVector.from_documents を使って行います。
この関数は、①DBの接続情報 ②格納したいドキュメント ③エンベディングモデル を引き渡すことだけで、ドキュメントのベクトル化とDBへの取り込みを行ってくれます。
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_postgres.vectorstores import PGVector
from langchain_community.document_loaders import PyPDFLoader
# 1. DBの接続情報
## Amazon Aurora PostgreSQL の接続情報を記入 ##
_host = "<データベースエンドポイント>" # TODO: 実行環境に合わせて修正
_user = "postgres"
_pass = "<設定したパスワード>" # TODO: 実行環境に合わせて修正
_port = "5432"
CONNECTION_STRING = f"postgresql+psycopg://{_user}:{_pass}@{_host}:{_port}"
# 2. 格納したいドキュメント
loader = PyPDFLoader(r"./tmp/aws-overview.pdf") # TODO: 実行環境に合わせて修正
pages = loader.load_and_split()
# 3. エンベディングモデル
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {"device": "cpu"}
encode_kwargs = {"normalize_embeddings": False}
embeddings = HuggingFaceEmbeddings(
model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs
)
# 1, 2, 3の情報を元にデータをベクトル化しDBに格納
PGVector.from_documents(
connection=CONNECTION_STRING, documents=pages, embedding=embeddings
)
準備が整ったら Code Editor インスタンスでプログラムを実行します。マシンスペックにもよりますが、2, 3分で完了します。
python vector_import.py
データ確認
プログラムが正常終了したらDBへ接続しベクトルデータが格納されているか確認します。
Code Editor インスタンスに postgresql がインストールされていない場合は、前回同様以下のコマンドでインストールを行います。
sudo apt update
sudo apt install postgresql postgresql-contrib
準備が出来たらログインします。
psql -h <データベースエンドポイント> -U <ユーザー名> -p <ポート番号>
PGVector.from_documents によって作成されたテーブルを確認します。
postgres=> \d
List of relations
Schema | Name | Type | Owner
--------+-------------------------+---------------+----------
public | langchain_pg_collection | table | postgres
public | langchain_pg_embedding | table | postgres
(2 rows)
langchain_pg_collection
と langchain_pg_embedding
が LangChain によって作成されたテーブルです。
これらテーブルの構成を見ていきましょう。
postgres=> \d langchain_pg_collection
Table "public.langchain_pg_collection"
Column | Type | Collation | Nullable | Default
-----------+-------------------+-----------+----------+---------
uuid | uuid | | not null |
name | character varying | | not null |
cmetadata | json | | |
Indexes:
"langchain_pg_collection_pkey" PRIMARY KEY, btree (uuid)
"langchain_pg_collection_name_key" UNIQUE CONSTRAINT, btree (name)
Referenced by:
TABLE "langchain_pg_embedding" CONSTRAINT "langchain_pg_embedding_collection_id_fkey" FOREIGN KEY (collection_id) REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE
postgres=> \d langchain_pg_embedding
Table "public.langchain_pg_embedding"
Column | Type | Collation | Nullable | Default
---------------+-------------------+-----------+----------+---------
id | character varying | | not null |
collection_id | uuid | | |
embedding | vector | | |
document | character varying | | |
cmetadata | jsonb | | |
Indexes:
"langchain_pg_embedding_pkey" PRIMARY KEY, btree (id)
"ix_cmetadata_gin" gin (cmetadata jsonb_path_ops)
"ix_langchain_pg_embedding_id" UNIQUE, btree (id)
Foreign-key constraints:
"langchain_pg_embedding_collection_id_fkey" FOREIGN KEY (collection_id) REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE
postgres=>
色々書いてありますが、要は以下のような構造です。
つまり、ベクトルデータそのものは langchain_pg_embedding
(図右)に格納され、それらをコレクションという単位で langchain_pg_collection
(図左)で管理しています。
仮にもう一度同じ方法でベクトルデータを格納した場合、新しく採番された uuid
が collection_id
としてエンベディングテーブルに追加されます。
では具体的な中身はというと、
以下のように embedding
カラムにベクトル化されたデータが格納されています。
また一番右端にある cmetadata
カラムを見ると分かる通り、これらデータはサンプルデータのAWSホワイトペーパーが1ページずつベクトル化されています。
id |collection_id |embedding |document |cmetadata |
-------------------------+--------------------+------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------+-----------------------------------------------+
26a36f35-<中略>-44b0-a733|1cabec0e-<中略>-40cf-|[0.0013618883,0.007772876<中略>6083,0.022092478,0.011488956,0.02407887,0|AWS ホワイトペーパー¶Amazon Web Services の概要¶Copy<中略>. © 2024 Amazon Web Services, Inc. and/or its affiliates |{"page": 0, "source": "./tmp/aws-overview.pdf"}|
5ec8628d-<中略>-46e6-95a1|1cabec0e-<中略>-40cf-|[0.03983145,-0.020729112,<中略>35172,0.020743001,0.0009571253,0.03718657|Amazon Web Services の概要 AWS ホワイトペーパー¶Amaz<中略>.b Services の概要: AWS ホワイトペーパー¶Copyright © 202使用|{"page": 1, "source": "./tmp/aws-overview.pdf"}|
33345333-<中略>-44a8-8b97|1cabec0e-<中略>-40cf-|[0.020863863,-0.04820327,<中略>26016,0.023011606,0.015820494,0.037092764|Amazon Web Services の概要 AWS ホワイトペーパー¶Tabl<中略>.Contents¶要約と序章 ..................................|{"page": 2, "source": "./tmp/aws-overview.pdf"}|
お疲れ様でした。実はベクトルデータの格納はこれにて完了です。
今回サンプルデータでPDFを用いましたが、もちろん csv や RDB 上の文字データなどでもベクトル化は可能です。その場合は vector_import.py
の LangChain 関数を変更して色々と試してみてください。
次項では上記データを用いてベクトルデータの近似計算を行ってみます。
近似計算
データをベクトル化した後に重要なのが、それをどう使うか、です。
例えば、ecサイトの検索精度を上げるため検索キーワードに対し適切な商品を提示する、などの目的があります。
今回は格納したAWSホワイトペーパーを使って幾つかのケースに沿って近似計算を行っていきましょう。
ケース1:キーワードに対して適切な参照ページを知りたい
これはホワイトペーパーに対するキーワード検索です。ただし単純な文字列一致検索ではなくベクトル検索です。よって打ち間違えたキーワード、または自然文であっても適切なページを検索することが出来ます。
データをインポートした Code Editor インスタンスからベクトル検索をしてみましょう。
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_postgres.vectorstores import PGVector
# 接続文字列の定義
_host = "<データベースエンドポイント>" # TODO: 実行環境に合わせて修正
_user = "postgres"
_pass = "<設定したパスワード>" # TODO: 実行環境に合わせて修正
_port = "5432"
CONNECTION_STRING = f"postgresql+psycopg://{_user}:{_pass}@{_host}:{_port}"
# Embeddingsの定義
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {"device": "cpu"}
encode_kwargs = {"normalize_embeddings": False}
embeddings = HuggingFaceEmbeddings(
model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs
)
# PostgreSQLの定義
db = PGVector.from_existing_index(connection=CONNECTION_STRING, embedding=embeddings)
# コサイン類似度の高い Top5 を検索
input_text = "sagemaker stadio" # あえて打ち間違えたキーワードで検索
result = db.similarity_search(query=input_text, k=5)
# 結果出力
for row in result:
page = row.metadata["page"]+1 # ページ数が0始まりで格納されているため1を足す
content = row.page_content[:100].replace("\n", " ") # 結果を見やすくするため先頭100文字だけ出力
print(f"page: {page} \t content: {content}")
以下のように出力されたら成功です。
※embedding 時に用いたモデルによってはヒットするページは異なります。
page: 95 content: Amazon Web Services の概要 AWS ホワイトペーパー Amazon SageMaker Autopilot Amazon SageMaker Autopilot は、データに基づい
page: 85 content: Amazon Web Services の概要 AWS ホワイトペーパー 各サービスは図の後に説明されています。ニーズに最も合ったサービスを判断するには、「 の選 択」を参照してください。 AWS 機
page: 96 content: Amazon Web Services の概要 AWS ホワイトペーパー Amazon SageMaker Edge Amazon SageMaker Edge は、 を最適化することで、エッジデバイ
page: 3 content: Amazon OpenSearch サーバーレス ...........................................................................
page: 98 content: Amazon Web Services の概要 AWS ホワイトペーパー ML コンピューティングインフラストラクチャを活用し、 SageMakerインフラストラクチャを 1 台 から数千台の に自動
4件目は目次のページでヒットしていますが、いずれのページも sagemaker を説明している適切なページが検索出来ていました。
もちろん、AWSホワイトペーパーのPDFそのものに "sagemaker stadio" と打ち間違えて検索した場合、これらのページはヒットしません。
一見地味ですが、例えば ecサイト上で0件ヒットとなりユーザーが離脱してしまうことの機会損失を考えると、打ち間違いを許容した検索というのは大変強力です。
ケース2:あるページに対して類似したページを知りたい
ちょっと今回のホワイトペーパーでは考えづらいケースですが、、
ページ毎に商品が紹介されているような商品カタログを想像してください。
その場合、ある商品に対して似た商品を抽出したい、というケースはあるのではないでしょうか。つまりレコメンデーションです。
「この商品を見ている人はこちらの商品も見ています」というのも、この近似計算を用いて実現することができます。
早速以下のスクリプトを実行してみましょう。
import psycopg2 as pg
import pandas as pd
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_postgres.vectorstores import PGVector
# 接続文字列の定義
_host = "<データベースエンドポイント>" # TODO: 実行環境に合わせて修正
_user = "postgres"
_pass = "<設定したパスワード>" # TODO: 実行環境に合わせて修正
_port = "5432"
con=pg.connect(host=_host, port=_port, user=_user, password=_pass)
# コサイン類似度の高い Top5 を検索
sql = """
SELECT
*,
1 - (e.embedding <=> sub.target_emb) AS cosine_similarity
FROM
langchain_pg_embedding e,
(
SELECT
embedding as target_emb
FROM
langchain_pg_embedding
WHERE
cmetadata = '{"page": 102, "source": "./tmp/aws-overview.pdf"}'
) sub
ORDER by cosine_similarity DESC
LIMIT 5
"""
result = pd.read_sql(sql=sql,con=con)
# 結果出力
for key, val in result[["cmetadata","cosine_similarity","document"]].iterrows():
page = val.cmetadata["page"]
similarity = val.cosine_similarity # コサイン類似度
content = val.document[:100].replace("\n", " ") # 結果を見やすくするため先頭100文字だけ出力
print(f"page: {page} \t similarity: {similarity} \t content: {content}")
このケースでは LangChain で提供されている関数 similarity_search を使わずに SQL だけで類似のページを抽出しています。これが pgvector の強みです。
上記のスクリプトでは "AWSホワイトペーパーの p.102" を条件に比較対象のページを指定しました (WHERE cmetadata = '{"page": 102, "source": "./tmp/aws-overview.pdf"}'
)。実務では商品の属性情報(カテゴリ、スペック、説明文)に置き換えることができ、例えば、あるスペックに合致した商品に対するレコメンド商品、なども簡単に抽出できます。みなさんも実務のデータに置き換えてイメージしてみてください。ここでお伝えしたいことは、シンプルな SQL だけで構造化情報を掛け合わせた近似計算が可能という事です。
さて、上記スクリプトの実行結果は以下のようになります。
page: 102 similarity: 1.0 content: Amazon Web Services の概要 AWS ホワイトペーパー AWS DeepRacer AWS DeepRacer は 1/18 スケールのレースカーで、強化学習 (RL) を始めるため
page: 97 similarity: 0.865711824213055 content: Amazon Web Services の概要 AWS ホワイトペーパー ML コンピューティングインフラストラクチャを活用し、 SageMakerインフラストラクチャを 1 台 から数千台の に自動
page: 98 similarity: 0.806267869098911 content: Amazon Web Services の概要 AWS ホワイトペーパー ができます。 MxNet で の使用を開始できます。AWS ML モデルを大規模に構築、トレーニング、 デプロイするためのプラ
page: 129 similarity: 0.7977725863456726 content: Amazon Web Services の概要 AWS ホワイトペーパー AWS Global Accelerator AWS Global Accelerator は、グローバルユーザーに提供するア
page: 22 similarity: 0.7969940183755184 content: Amazon Web Services の概要 AWS ホワイトペーパー し、監査ログを保持してコンプライアンスとアクティビティのレポートを可能にする FinSpace よう に強制します。 には、分
1件目の "similarity" が 1.0 になっています。これはどういうことでしょうか。
page 番号を見てみると、今回条件で使ったp.102であることが分かります。つまりこれはこのレコメンデーションにおける比較元です(共起元とも言います)。自分自身を比較すればもちろん完全一致しますので、類似度は 1.0 という事です。
レコメンデーションでは比較元に対して似ているアイテムを出したいので、この共起元を除いた結果(共起先とも言います)をお客様へ提示することになります。なお実務でこれを実装するとなると、類似度いくつまでの結果を提示すべきか?出す順番は類似度順のままでいいのか?など検討すべき事項はまだまだあります。ただこれはこれで別の記事が書けるネタなのでまたの機会にします。
ケース3:検索結果を元に要約してほしい...(次回)
これは検索結果と LLM を繋げることで実現し、RAG や AIエージェント と呼ばれる仕組みとなります。次回はこの仕組みを簡易アプリとして実装することで pgvector入門の締めくくりとしたいと思います。
おわりに
前回の環境構築編と比べると pgvector の使い方がより具体的になったのではないでしょうか。記事の途中でも少し触れましたが、ベクトル化に使用するモデルを変えることで結果が異なってきます。Hugging Face 等を通じて様々なモデルが無料で利用できますので、目的やデータに合わせて色々なモデルを試してみてください。
Discussion