Closed7
pgvectorを試してみる
kun432
fly.ioだけでベクトル検索を使ったAPI作るのが完結しそうという意味で興味がある。
kun432こういうのもあるらしい。20x fasterだそうな。
pgvectorちゃんと試してからやってみる
kun432参考）
コンテナ起動
$ docker run -it --rm \
--name pgvector-test \
-e POSTGRES_PASSWORD=password \
-e POSTGRES_INITDB_ARGS="--encoding=UTF8 --no-locale" \
-e TZ=Asia/Tokyo \
-v data:/var/lib/postgresql/data \
-p 5432:5432 \
-d \
postgres:16
コンテナにログイン
$ docker exec -ti pgvector-test /bin/bash
プロセス確認
$ apt-get update
$ apt-get install -y procps
$ ps auxw | grep postgres
postgres 1 0.0 0.0 220056 29184 pts/0 Ss+ 14:06 0:00 postgres
postgres 27 0.0 0.0 220192 9064 ? Ss 14:06 0:00 postgres: checkpointer
postgres 28 0.0 0.0 220208 6504 ? Ss 14:06 0:00 postgres: background writer
postgres 30 0.0 0.0 220056 9832 ? Ss 14:06 0:00 postgres: walwriter
postgres 31 0.0 0.0 221652 8552 ? Ss 14:06 0:00 postgres: autovacuum launcher
postgres 32 0.0 0.0 221636 7784 ? Ss 14:06 0:00 postgres: logical replication launcher
root 4645 0.0 0.0 6332 1792 pts/1 S+ 14:12 0:00 grep postgres
DBに接続
$ psql -h localhost -U postgres
psql (16.2 (Debian 16.2-1.pgdg120+2))
Type "help" for help.
postgres=#
ロケールやタイムゾーン、現時点での拡張を確認。
postgres=# \l
List of databases
Name | Owner | Encoding | Locale Provider | Collate | Ctype | ICU Locale | ICU Rules | Access privileges
-----------+----------+----------+-----------------+---------+-------+------------+-----------+-----------------------
postgres | postgres | UTF8 | libc | C | C | | |
template0 | postgres | UTF8 | libc | C | C | | | =c/postgres +
| | | | | | | | postgres=CTc/postgres
template1 | postgres | UTF8 | libc | C | C | | | =c/postgres +
| | | | | | | | postgres=CTc/postgres
(3 rows)
postgres=# select NOW();
now
------------------------------
2024-04-24 14:13:41.51503+09
(1 row)
postgres=# \dx
List of installed extensions
Name | Version | Schema | Description
---------+---------+------------+------------------------------
plpgsql | 1.0 | pg_catalog | PL/pgSQL procedural language
(1 row)
postgres=# \q
ではpgvectorのビルドとインストール。
依存パッケージのインストール。
$ apt-get install -y git gcc make postgresql-server-dev-16
pgvectorをクローンしてビルド・インストール。
$ cd /tmp
$ git clone --branch v0.6.2 https://github.com/pgvector/pgvector.git && cd pgvector
$ make
$ make install
pgvectorの有効化。
$ psql -h localhost -U postgres
psql (16.2 (Debian 16.2-1.pgdg120+2))
Type "help" for help.
postgres=# \dx
List of installed extensions
Name | Version | Schema | Description
---------+---------+------------+------------------------------
plpgsql | 1.0 | pg_catalog | PL/pgSQL procedural language
(1 row)
postgres=# CREATE EXTENSION vector;
CREATE EXTENSION
postgres=# \dx
List of installed extensions
Name | Version | Schema | Description
---------+---------+------------+------------------------------------------------------
plpgsql | 1.0 | pg_catalog | PL/pgSQL procedural language
vector | 0.6.2 | public | vector data type and ivfflat and hnsw access methods
(2 rows)
公式READMEに従ってテーブル作成＋データ投入。
postgres=# CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));
CREATE TABLE
postgres=# INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,9]'), ('[0,1,2]'), ('[3,4,5]'), ('[6,7,8]'), ('[9,0,1]'), ('[2,3,4]'), ('[5,6,7]'), ('[8,9,0]');
INSERT 0 10
L2距離で検索
postgres=# SELECT *, embedding <-> '[3,1,2]' AS l2_distance FROM items ORDER BY l2_distance LIMIT 5;
id | embedding | l2_distance
----+-----------+-------------------
1 | [1,2,3] | 2.449489742783178
4 | [0,1,2] | 3
8 | [2,3,4] | 3
5 | [3,4,5] | 4.242640687119285
2 | [4,5,6] | 5.744562646538029
(5 rows)
ドット積で検索
postgres=# SELECT *, (embedding <#> '[3,1,2]') * -1 AS inner_product FROM items ORDER BY inner_product LIMIT 5;
id | embedding | inner_product
----+-----------+---------------
4 | [0,1,2] | 5
1 | [1,2,3] | 11
8 | [2,3,4] | 17
5 | [3,4,5] | 23
7 | [9,0,1] | 29
(5 rows)
コサイン類似度で検索
postgres=# SELECT *, 1 - (embedding <=> '[3,1,2]') AS cosine_similarity FROM items ORDER BY cosine_similarity LIMIT 5;
id | embedding | cosine_similarity
----+-----------+--------------------
4 | [0,1,2] | 0.5976143046671968
10 | [8,9,0] | 0.7324296566704842
1 | [1,2,3] | 0.7857142857142857
8 | [2,3,4] | 0.8436958338752907
7 | [9,0,1] | 0.8559079373463852
(5 rows)
インデックスを作成する。HNSW・コサイン類似度の場合。
postgres=# CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
CREATE INDEX
postgres=# SELECT *, 1 - (embedding <=> '[3,1,2]') AS cosine_similarity FROM items ORDER BY cosine_similarity LIMIT 5;
id | embedding | cosine_similarity
----+-----------+--------------------
4 | [0,1,2] | 0.5976143046671968
10 | [8,9,0] | 0.7324296566704842
1 | [1,2,3] | 0.7857142857142857
8 | [2,3,4] | 0.8436958338752907
7 | [9,0,1] | 0.8559079373463852
(5 rows)
LangChainから繋いでみる。ローカルのJupyterLabから。
パッケージインストール
!pip install --upgrade --quiet langchain langchain-core langchain-openai langchain_postgres langchain-text-splitters python-dotenv grandalf
OpenAI APIキー読み込み。.envは予めセットしておく。
from dotenv import load_dotenv
load_dotenv(verbose=True)
ドキュメント準備。ここは過去のやり方を踏襲。
from pathlib import Path
import requests
import re
def replace_heading(match):
level = len(match.group(1))
return '#' * level + ' ' + match.group(2).strip()
# Wikipediaからのデータ読み込み
wiki_titles = ["オグリキャップ"]
for title in wiki_titles:
response = requests.get(
"https://ja.wikipedia.org/w/api.php",
params={
"action": "query",
"format": "json",
"titles": title,
"prop": "extracts",
# 'exintro': True,
"explaintext": True,
},
).json()
page = next(iter(response["query"]["pages"].values()))
wiki_text = f"# {title}\n\n## 概要\n\n"
wiki_text += page["extract"]
wiki_text = re.sub(r"(=+)([^=]+)\1", replace_heading, wiki_text)
wiki_text = re.sub(r"\t+", "", wiki_text)
wiki_text = re.sub(r"\n{3,}", "\n\n", wiki_text)
data_path = Path("data")
if not data_path.exists():
Path.mkdir(data_path)
# markdown(.md)ファイルとして出力
with open(data_path / f"{title}.md", "w") as fp:
fp.write(wiki_text)
import glob
import os
from langchain_text_splitters import MarkdownHeaderTextSplitter
from langchain_core.documents import Document
def text_split(text, max_length=400):
"""
"""
chunks = re.split(r'(?<=[。！？\n])', text)
chunks = [s for s in chunks if s.strip()]
temp_chunk = ""
final_chunks = []
for chunk in chunks:
if len(temp_chunk + chunk) <= max_length:
temp_chunk += chunk
else:
final_chunks.append(temp_chunk)
temp_chunk = chunk
if temp_chunk:
final_chunks.append(temp_chunk)
return final_chunks
sections_for_delete = ["競走成績", "外部リンク", "参考文献"]
headers_to_split_on = [
("#", "Header 1"),
("##", "Header 2"),
("###", "Header 3"),
("####", "Header 4"),
("#####", "Header 5"),
("######", "Header 6"),
]
files = glob.glob('data/*.md')
splits = []
for file in files:
with open(file) as f:
md = f.read()
markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
docs_split = markdown_splitter.split_text(md)
docs_for_delete = []
for idx, d in enumerate(docs_split):
metadatas = []
header_keys = []
d.metadata["source"] = file
for m in d.metadata:
if m.startswith("Header"):
metadatas.append(d.metadata[m])
header_keys.append(m)
# 削除対象のセクションを含むドキュメントを後で削除するためにそのインデックス登録しておく
if d.metadata[m] in sections_for_delete:
docs_for_delete.append(idx)
# セクションの階層を結合、パンくずリストとしてセクション情報に追加
if len(metadatas) > 0:
d.metadata["section"] = metadata_str = " > ".join(metadatas)
for k in header_keys:
if k.startswith("Header"):
del d.metadata[k]
# 削除対象セクションの削除
docs = [item for i, item in enumerate(docs_split) if i not in docs_for_delete]
for d in docs:
chunks = text_split(d.page_content, 500)
if len(chunks) == 1:
splits.append(d)
else:
for idx, chunk in enumerate(chunks, start=1):
metadata = d.metadata.copy()
metadata["section"] += f"（{idx}）"
splits.append(Document(page_content=chunk, metadata=metadata))
for i in splits[:5]:
print(i.metadata)
print(i.page_content[:60] + "...")
print("====")
for i in splits[:5]:
print(i.metadata)
print(i.page_content[:60] + "...")
print("====")
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 概要'}
オグリキャップ（欧字名:Oguri Cap、1985年3月27日 - 2010年7月3日）は、日本の競走馬、種牡馬。
1...
====
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > デビューまで > 誕生に至る経緯'}
オグリキャップの母・ホワイトナルビーは競走馬時代に馬主の小栗孝一が所有し、笠松競馬場の調教師鷲見昌勇が管理した。ホワイト...
====
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > デビューまで > 誕生・生い立ち > 稲葉牧場時代'}
オグリキャップは1985年3月27日の深夜に誕生した。誕生時には右前脚が大きく外向しており、出生直後はなかなか自力で立ち...
====
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > デビューまで > 誕生・生い立ち > 美山育成牧場時代'}
1986年の10月、ハツラツは岐阜県山県郡美山町（現：山県市）にあった美山育成牧場に移り、3か月間馴致を施された。当時の...
====
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 競走馬時代 > 笠松競馬時代 > 競走内容（1）'}
1987年1月28日に笠松競馬場の鷲見昌勇厩舎に入厩。登録馬名は「オグリキヤツプ」。ダート800mで行われた能力試験を5...
====
LangChainからpgvectorに接続してドキュメントを追加。
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_postgres import PGVector
from langchain_postgres.vectorstores import PGVector
connection = "postgresql+psycopg://postgres:password@localhost:5432/postgres"
collection_name = "oguricap"
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = PGVector(
embeddings=embeddings,
collection_name=collection_name,
connection=connection,
use_jsonb=True,
)
vectorstore.add_documents(splits, ids=[idx for idx, doc in enumerate(splits, start=1)])
検索
search_results = vectorstore.similarity_search_with_score("オグリキャップの血統は？", k=5)
for r in search_results:
document = r[0]
score = r[1]
print(score)
print(document.metadata)
print(document.page_content[:60] + "...")
print("====")
0.43296221298204485
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 概要'}
オグリキャップ（欧字名:Oguri Cap、1985年3月27日 - 2010年7月3日）は、日本の競走馬、種牡馬。
1...
====
0.44227665662765503
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 特徴・評価 > 競走馬名および愛称・呼称'}
競走馬名「オグリキャップ」の由来は、馬主の小栗が使用していた冠名「オグリ」に父ダンシングキャップの馬名の一部「キャップ」...
====
0.49260693876660244
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 特徴・評価 > 走行・レースぶりに関する特徴・評価（2）'}
オグリキャップは肢のキック力が強く、瞬発力の強さは一回の蹴りで前肢を目いっぱいに延ばし、浮くように跳びながら走るため、こ...
====
0.5119419395923615
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 特徴・評価 > 身体面に関する特徴・評価（2）'}
オグリキャップの体力面について、競馬関係者からは故障しにくい点や故障から立ち直るタフさを評価する声が挙がっている。輸送時...
====
0.512657598038244
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 人気 > 概要（2）'}
お笑い芸人の明石家さんまは雑誌『サラブレッドグランプリ』のインタビューにおいて、オグリキャップについて「マル地馬で血統も...
====
retrieverで検索
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
search_result = retriever.invoke("オグリキャップの血統は？")
for r in search_result:
print(r.metadata)
print(r.page_content[:60] + "...")
print("====")
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 概要'}
オグリキャップ（欧字名:Oguri Cap、1985年3月27日 - 2010年7月3日）は、日本の競走馬、種牡馬。
1...
====
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 特徴・評価 > 競走馬名および愛称・呼称'}
競走馬名「オグリキャップ」の由来は、馬主の小栗が使用していた冠名「オグリ」に父ダンシングキャップの馬名の一部「キャップ」...
====
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 特徴・評価 > 走行・レースぶりに関する特徴・評価（2）'}
オグリキャップは肢のキック力が強く、瞬発力の強さは一回の蹴りで前肢を目いっぱいに延ばし、浮くように跳びながら走るため、こ...
====
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 特徴・評価 > 身体面に関する特徴・評価（2）'}
オグリキャップの体力面について、競馬関係者からは故障しにくい点や故障から立ち直るタフさを評価する声が挙がっている。輸送時...
====
{'source': 'data/オグリキャップ.md', 'section': 'オグリキャップ > 人気 > 概要（2）'}
お笑い芸人の明石家さんまは雑誌『サラブレッドグランプリ』のインタビューにおいて、オグリキャップについて「マル地馬で血統も...
====
データベース側でどうなっているか確認してみる。
$ psql -h localhost -U postgres
新しくlangchain_pg_collectionとlangchain_pg_embeddingが作成されている。
postgres=# \d
List of relations
Schema | Name | Type | Owner
--------+-------------------------+----------+----------
public | items | table | postgres
public | items_id_seq | sequence | postgres
public | langchain_pg_collection | table | postgres
public | langchain_pg_embedding | table | postgres
(4 rows)
postgres=# \d langchain_pg_collection
uuid | uuid | | not null |
name | character varying | | not null |
cmetadata | json | | |
postgres=# SELECT * from langchain_pg_collection;
uuid | name | cmetadata
--------------------------------------+----------+-----------
4baa47ed-2f00-4585-b331-b0ea4a13d1b3 | oguricap | null
(1 row)
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=# \x
Expanded display is on.
postgres=# SELECT id, collection_id, document, cmetadata FROM langchain_pg_embedding WHERE id='1';
-[ RECORD 1 ]-+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
id | 1
collection_id | 4baa47ed-2f00-4585-b331-b0ea4a13d1b3
document | オグリキャップ（欧字名:Oguri Cap、1985年3月27日 - 2010年7月3日）は、日本の競走馬、種牡馬。(snip)
cmetadata | {"source": "data/オグリキャップ.md", "section": "オグリキャップ > 概要"}
なるほど、LangChain専用って感じになるのね。あと、ドキュメントを見ると、
現在のところ、スキーマ変更時のデータ移行を簡単にサポートする仕組みはありません。そのため、ベクターストアのスキーマを変更した場合、ユーザーはテーブルを再作成し、ドキュメントを追加し直す必要があります。これが気になる場合は、別のベクターストアを使用してください。そうでない場合は、この実装で問題ないでしょう。
このあたり見てる限り、まあしょうがないとは言え、運用的にはいろいろ気になるところよなぁ。
kun432LlamaIndexは試さないけど、ドキュメント見る限りは細かいところまで設定とかできそう。とはいえLlamaIndexのドキュメントのメタデータはかなり複雑な気がするので、それはそれで面倒かなという気はする。
kun432削除。--rmつけてるので落としたら削除される。
$ docker stop pgvector-test
PostgreSQLを使うメリットは、ベクトルDB以外に管理したいデータが必要な場合になると思う。例えば、
- ベクトルDBに登録するドキュメントを顧客IDごとに紐づけておいて、管理画面でCRUDできるようにする。
- ドキュメントのテキスト更新等があった場合に、ベクトルDBも更新する。
- 上記以外にもDB管理したいデータがある。
というようなケース。このあたりの運用を考えると、ネイティブできちんと設計してやるのが色んな意味で良さそうではあるかな。
あくまでも個人的にだけど、
- あまりRDBはがっつり触れてきていない
- テキストの更新とベクトルデータの更新をアトミックにやれないのではないか？（と思っている）
というところでやや面倒だなと。
これならば、モジュール使ってサーバサイドでベクトル生成してくれて、アトミックにCRUDできるWeaviateとかがいいなあと思ってしまう。RDBとは別にしたほうがいろいろ身軽になれそうな気もしているし。
あとはQdrant meilisearchあたりもモジュール的な仕組みができていたはずだし、多分この辺の流れは今後も増えるんじゃないかと思う。meilisearchなら全文検索寄りになりそうでそれはそれで今後必要になりそうだし。
まあPostgreSQLのスキルセットがあるならこちらのほうがいいかもね。
このスクラップは11日前にクローズされました