ベクトルDBとして有名なPineconeのベクトル検索を試してみました。
今回はメタデータを使ったフィルタリングも使ってみました。
なおエンベディングにはCohereのモデルを使い、実装にはLangChainを活用しています。
事前準備
必要なPythonパッケージをインストールします。
pip install pinecone-client langchain
Pineconeのサイトにて無料でサインアップし、APIキーを取得します。
https://www.pinecone.io/
Cohereを使う場合は下記サイトにてアカウントを作成し、トライアルのAPIキーを取得します。
https://cohere.com/
取得した各APIキーを .env ファイルに保存し、Pythonコードと同じディレクトリに配置します。
インデックス作成とデータ追加
下記コードは新しいインデックスを作成し、3件のデータを追加します。
サンプルデータはWikipediaから引用した文章になります。
ちなみにPineconeで言うインデックスはOracle DBで言うスキーマのような位置づけみたいです。
import os
from dotenv import load_dotenv
load_dotenv()
from pinecone import Pinecone, PodSpec
from langchain.embeddings import CohereEmbeddings
import json
# embedding設定値
COHERE_API_KEY = os.environ['COHERE_API_KEY']
embeddings = CohereEmbeddings(cohere_api_key=COHERE_API_KEY, model="embed-multilingual-v3.0")
# pinecone設定値
PINECONE_API_KEY = os.environ['PINECONE_API_KEY']
pc = Pinecone(api_key=PINECONE_API_KEY)
# index作成
index_name = "sample-index"
pc.create_index(
name=index_name,
dimension=1024,
metric="cosine",
spec=PodSpec(
environment="gcp-starter"
)
)
index = pc.Index(index_name)
# サンプルデータとしてWikipediaから引用した文章を用意 (メタデータにカテゴリ情報を付与)
sample_texts = [{"text": "生成的人工知能モデルは、訓練データの規則性や構造を訓練において学習することで、訓練データに含まれない新しいデータを生成することができる[8][9]。ジェネレーティブAI、ジェネラティブAIともよばれる。", "category": "AI"}, {"text": "富士山（ふじさん）は、静岡県（富士宮市、富士市、裾野市、御殿場市、駿東郡小山町）と山梨県（富士吉田市、南都留郡鳴沢村）に跨る活火山である。", "category": "mountain"}, {"text": "『鬼滅の刃』は、吾峠呼世晴による日本の漫画作品。", "category": "comics"}]
# embedding
for i in range(len(sample_texts)):
# Cohereの embed-multilingual-v3.0 でエンベディング
vector = embeddings.embed_query(sample_texts[i]["text"])
# エンベディング結果をPineconeのインデックスに保存
index.upsert(
vectors=[
{"id":str(i),
"values": vector,
"metadata": {"text": sample_texts[i]["text"], "category": sample_texts[i]["category"]}
}],
namespace='sample-namespace'
)
ベクトル検索
｢観光地」という単語でベクトル検索するコードです。
import os
from dotenv import load_dotenv
load_dotenv()
from pinecone import Pinecone, PodSpec
from langchain.embeddings import CohereEmbeddings
import json
# embedding設定値
COHERE_API_KEY = os.environ['COHERE_API_KEY']
embeddings = CohereEmbeddings(cohere_api_key=COHERE_API_KEY, model="embed-multilingual-v3.0")
# pinecone設定値
PINECONE_API_KEY = os.environ['PINECONE_API_KEY']
pc = Pinecone(api_key=PINECONE_API_KEY)
# 作成済みインデックスを指定
index_name = "sample-index"
index = pc.Index(index_name)
# ベクトル検索
query_text = "観光地"
query_embed = embeddings.embed_query(query_text)
res = index.query(
vector=query_embed,
top_k=3,
include_metadata=True,
namespace="sample-namespace"
)
print(res)
結果は以下の通りです。
{'matches': [{'id': '1',
'metadata': {'category': 'mountain',
'text': '富士山（ふじさん）は、静岡県（富士宮市、富士市、裾野市、御殿場市、駿東郡小山町）と山梨県（富士吉田市、南都留郡鳴沢村）に跨る活火山である。'},
'score': 0.541224539,
'values': []},
{'id': '2',
'metadata': {'category': 'comics',
'text': '『鬼滅の刃』は、吾峠呼世晴による日本の漫画作品。'},
'score': 0.421164662,
'values': []},
{'id': '0',
'metadata': {'category': 'AI',
'text': '生成的人工知能モデルは、訓練データの規則性や構造を訓練において学習することで、訓練データに含まれない新しいデータを生成することができる[8][9]。ジェネレーティブAI、ジェネラティブAIともよばれる。'},
'score': 0.306574225,
'values': []}],
'namespace': 'sample-namespace',
'usage': {'read_units': 6}}
観光地と一番関連のある富士山の紹介文がトップに来ています。
またメタデータを使ったフィルタリングも出来るということで、「AI」カテゴリに絞ってベクトル検索してみます。
検索ワードは「機械学習」にしました。
import os
from dotenv import load_dotenv
load_dotenv()
from pinecone import Pinecone, PodSpec
from langchain.embeddings import CohereEmbeddings
import json
# embedding設定値
COHERE_API_KEY = os.environ['COHERE_API_KEY']
embeddings = CohereEmbeddings(cohere_api_key=COHERE_API_KEY, model="embed-multilingual-v3.0")
# pinecone設定値
PINECONE_API_KEY = os.environ['PINECONE_API_KEY']
pc = Pinecone(api_key=PINECONE_API_KEY)
# 作成済みインデックスを指定
index_name = "sample-index"
index = pc.Index(index_name)
# ベクトル検索
query_text = "機械学習"
query_embed = embeddings.embed_query(query_text)
# filterにメタデータのフィルタリング情報を設定
res = index.query(
vector=query_embed,
filter={
"category": {"$eq": "AI"}
},
top_k=3,
include_metadata=True,
namespace="sample-namespace"
)
print(res)
結果は以下の通りです。
{'matches': [{'id': '0',
'metadata': {'category': 'AI',
'text': '生成的人工知能モデルは、訓練データの規則性や構造を訓練において学習することで、訓練データに含まれない新しいデータを生成することができる[8][9]。ジェネレーティブAI、ジェネラティブAIともよばれる。'},
'score': 0.646658897,
'values': []}],
'namespace': 'sample-namespace',
'usage': {'read_units': 6}}
ちゃんとメタデータの「category」が「AI」の文章だけ検索対象になりました。
簡単ですが以上になります。
少しだけ使ってみて、非常に手軽な印象を受けました。
Slowly But Surely 