ベクトル特化型データベースサービス「Pinecone」でセマンティック・キーワード検索をやってみた
本記事では、マネージド・ベクトル・データベースの「Pinecone」を活用して、セマンティック・キーワード検索を実施していきます。
ベクトル分析は、類似性の計算やレコメンドの作成などで使われる機械学習のメジャーな分析手法ですが、それに特化したユニークなデータベースのSaaSを見つけちゃいました。
Pineconeについて
Pineconeは2019年にカリフォルニア州で創業されたスタートアップです。もともとAmazon SageMakerの開発に携わっていた方が創業したようで、シード期ながら$10Mもの資金調達に成功しており、期待値の高さが伺えます。GoogleやPinterestのようなビッグカンパニーが使用している、高次元ベクトルのデータを格納できる機械学習用のデータベースを一般企業にも広めたい、というモチベーションのもとPineconeを開発・提供しています。
Pinecone lands $10M seed for purpose-built machine learning database | TechCrunch
Pineconeは、数十億のベクトルを10msで検索可能なスケーラビリティとパフォーマンスを提供し、本番のアプリケーションにすぐ導入できるインテグレーション性の高さが強みです。マネージドサービスなので、運用や監視のコストも最小限。主なユースケースとしては、以下の5点が挙げられています。
- Semantic Search: セマンティック検索
- Unstructured Data Search: 非構造化データ検索（画像や音声など）
- Deduplication and Record Matching: 重複排除とレコードマッチング
- Recommendations and Ranking: レコメンデーションとランキング
- Detection and Classfication: 検知と分類
本記事では、Pineconeの無料アカウントを作成し、QuickstartとHybrid Searchのチュートリアルをやっていきます。
Quickstart
まずはアカウントを作成していきます。Pineconeの公式サイトにてStart for Freeをクリックします。
Pineconeは最初からGoogle認証でアカウント作成できます！これはありがたい。
認証が通るとPineconeのホーム画面に到着します。ただ、Pineconeは主にJupyter Notebookなどから操作する製品なので、あまりWeb UIは使用しません。
それでは、Quickstartに沿って動作確認していきます。環境としてhello_pinecone.ipynb - Colaboratoryが提供されているので、こちらを開きます。まずはpinecone-clientをインストールします。
pip install -qU pip pinecone-client
PineconeのUIからAPI Keyをコピーし、pinecone.init()の引数に渡して実行します。
import pinecone pinecone.init(api_key="XXXXX-XXXXXXX-XXXXXX-XXXXXX", environment='us-west1-gcp')
セットアップ後、PineconeのIndexを作成します。PineconeでのIndexは「ベクトルデータの一番上位の階層単位」として定義されています。このIndexがベクトルデータを保存するためのキーになったり、クエリを発行する時の起点になったりします。
dimensions = 3 pinecone.create_index(name=index_name, dimension=dimensions, metric="cosine", shards=1)
pinecone.list_indexes()でも確認可能ですが、せっかくなのでUI側からIndexが作成されているかどうか確認してみます。
次に、作成したIndexにデータを流していきます。サンプルとして2レコードだけ挿入します。
index = pinecone.Index(index_name=index_name)
import pandas as pd
df = pd.DataFrame(
data={
"id": ["A", "B"],
"vector": [[1., 1., 1.], [1., 2., 3.]]
})
index.upsert(vectors=zip(df.id, df.vector)) # insert vectors
作成したベクトルをクエリするには、以下の関数を実行します。
index.query(
queries=[[2., 2., 2.], [2., 4., 6.]],
top_k=5,
include_values=True) # returns top_k matches for every query
無事レスポンスが返ってきました。使い心地はよくあるNoSQLといった具合。
{'results': [{'matches': [{'id': 'A',
'score': 0.99999994,
'values': [0.99999994, 0.99999994, 0.99999994]},
{'id': 'B',
'score': 0.925820053,
'values': [1.0, 2.0, 3.0]}],
'namespace': ''},
{'matches': [{'id': 'B',
'score': 0.999999881,
'values': [1.0, 2.0, 3.0]},
{'id': 'A',
'score': 0.925820053,
'values': [0.99999994, 0.99999994, 0.99999994]}],
'namespace': ''}]}
無料プランの場合は1 Indexしか作成できないので、最後にIndexを削除しておきましょう。
pinecone.delete_index(index_name)
以上で動作確認は完了です。次節より、もう少し具体的な分析を試していきます。
Hybrid Search
Pineconeでは様々なベクトル分析をサポートしていますが、中でも、セマンティック検索にトラディショナルなキーワード検索が組み合わせ使えるHybrid Searchが特徴的です。セマンティック検索では、文字よりもその意味を解釈して検索をかけることができますが、文字自体も同時にフィルタリング対象としたい、という場合に便利な機能です。
早速、Basic Hybrid Search in Pineconeに沿って試していきます。このチュートリアルもNotebookが用意されているので活用していきます。
hybrid_search.ipynb - Colaboratory
まず、サンプルデータとして以下の10行を用意します。
all_sentences = [
"purple is the best city in the forest",
"No way chimps go bananas for snacks!",
"it is not often you find soggy bananas on the street",
"green should have smelled more tranquil but somehow it just tasted rotten",
"joyce enjoyed eating pancakes with ketchup",
"throwing bananas on to the street is not art",
"as the asteroid hurtled toward earth becky was upset her dentist appointment had been canceled",
"I'm getting way too old. I don't even buy green bananas anymore.",
"to get your way you must not bombard the road with yellow fruit",
"Time flies like an arrow; fruit flies like a banana"
]
セマンティック検索では、sentence-transformersライブラリの学習済みモデルを使用していきます。
from sentence_transformers import SentenceTransformer
model = SentenceTransformer('flax-sentence-embeddings/all_datasets_v3_mpnet-base')
all_embeddings = model.encode(all_sentences)
all_embeddings.shape
all_embeddings.shapeの出力結果から、all_sentencesは埋め込みが10、次元が768ということが読み取れます。続いて、文のトークン化に使用するHuggingFaceのライブラリ、transfo-xl-wt103をインストールします。
from transformers import AutoTokenizer
# transfo-xl tokenizer uses word-level encodings
tokenizer = AutoTokenizer.from_pretrained('transfo-xl-wt103')
all_tokens = [tokenizer.tokenize(sentence.lower()) for sentence in all_sentences]
all_tokens[0]
ライブラリのセットアップは以上です。先ほどと同様、API KeyでPineconeと接続しIndexを作成します。
import pinecone
pinecone.init(api_key='XXXXX-XXXXXXX-XXXXXX-XXXXXX', environment='us-west1-gcp')
pinecone.list_indexes() # check if keyword-search index already exists
pinecone.create_index(name='keyword-search', dimension=all_embeddings.shape[1])
index = pinecone.Index('keyword-search')
作成したIndexに対してall_embeddingsとall_tokensをInsertしていきます。
upserts = []
for i, (embedding, tokens) in enumerate(zip(all_embeddings, all_tokens)):
upserts.append((str(i), embedding.tolist(), {'tokens': tokens}))
# then we upsert
index.upsert(vectors=upserts)
格納が完了したので、さっそくクエリを投げていきましょう。まずはセマンティック検索単体で試してみます。エンコードしたクエリ文を引数に渡し、クエリ結果の上位10件を表示させます。
query_sentence = "there is an art to getting your way and throwing bananas on to the street is not it" xq = model.encode([query_sentence]).tolist() result = index.query(xq, top_k=5, includeMetadata=True) result
結果は類似度スコアが高い順に並び、関連性の高いトークン情報も含まれて返ってきます。
{"results": [{
"matches": [
{
"id": "5",
"metadata": {
"tokens": [
"throwing", "bananas", "on", "to", "the", "street", "is", "not", "art"
]
},
"score": 0.732851744,
"values": []
},
{
"id": "8",
"metadata": {
"tokens": [
"to", "get", "your", "way", "you", "must", "not", "bombard", "the", "road", "with", "yellow", "fruit"
]
},
"score": 0.574426889,
"values": []
},
{
"id": "2",
"metadata": {
"tokens": [
"it", "is", "not", "often", "you", "find", "soggy", "bananas", "on", "the", "street"
]
},
"score": 0.500876784,
"values": []
},
{
"id": "1",
"metadata": {
"tokens": [
"no", "way", "chimps", "go", "bananas", "for", "snacks", "!"
]
},
"score": 0.376693517,
"values": []
},
{
"id": "9",
"metadata": {
"tokens": [
"time", "flies", "like", "an", "arrow", ";", "fruit", "flies", "like", "a", "banana"
]
},
"score": 0.338697314,
"values": []
},
{
"id": "7",
"metadata": {
"tokens": [
"i", "'m", "getting", "way", "too", "old.", "i", "don", "'t", "even", "buy", "green", "bananas", "anymore", "."
]
},
"score": 0.324042052,
"values": []
},
{
"id": "0",
"metadata": {
"tokens": [
"purple", "is", "the", "best", "city", "in", "the", "forest"
]
},
"score": 0.145487517,
"values": []
},
{
"id": "3",
"metadata": {
"tokens": [
"green", "should", "have", "smelled", "more", "tranquil", "but", "somehow", "it", "just", "tasted", "rotten"
]
},
"score": 0.137328982,
"values": []
},
{
"id": "4",
"metadata": {
"tokens": [
"joyce", "enjoyed", "eating", "pancakes", "with", "ketchup"
]
},
"score": 0.0915388912,
"values": []
},
{
"id": "6",
"metadata": {
"tokens": [
"as", "the", "asteroid", "hurtled", "toward", "earth", "becky", "was", "upset", "her", "dentist", "appointment", "had", "been", "canceled"
]
},
"score": -0.0585537851,
"values": []
}
],
"namespace": ""
}]}
次に、このトークン情報を活用して、キーワード検索も絡めて実行します。クエリを実行する際にトークンにbananasが含まれているレコードでフィルタリングするよう指定します。
result = index.query(xq, top_k=10, filter={'tokens': 'bananas'})
ids = [int(x['id']) for x in result['results'][0]['matches']]
for i in ids:
print(all_sentences[i])
結果、4行がヒットしました。これがセマンティック検索とキーワード検索のハイブリット型、というわけですね。
throwing bananas on to the street is not art it is not often you find soggy bananas on the street No way chimps go bananas for snacks! I'm getting way too old. I don't even buy green bananas anymore.
このキーワード検索のフィルタリングはORやAND、IN条件でも可能です。
result = index.query(xq, top_k=10, filter={'$or': [
{'tokens': 'bananas'},
{'tokens': 'way'}
]})
ids = [int(x['id']) for x in result['results'][0]['matches']]
for i in ids:
print(all_sentences[i])
bananasかwayでヒットした検索結果です。$orオペレータの他にも、$andや$nin（not in）、$ne（not equal）もあるので、大体の条件検索は可能です。
throwing bananas on to the street is not art to get your way you must not bombard the road with yellow fruit it is not often you find soggy bananas on the street No way chimps go bananas for snacks! I'm getting way too old. I don't even buy green bananas anymore.
このようにPineconeでは、Indexとベクトルで構成されたデータセットに対して、機械学習的なフィルタリングをかけながらクエリができるのが特徴みたいですね。チュートリアルの実践は以上です。
所感
使い勝手がよく、アプローチがユニークで良い製品だなという印象を持ちました。私自身は機械学習の専門家ではないので、ぜひその手のエンジニアの方にも試していただいて感想をお聞きしてみたいです。Pineconeを使うと、普通のKey-Value型のデータベースとはどうパフォーマンスが異なるのかも気になります。今後、検証記事や事例が増えてくるのを楽しみにしてます。
本アドベントカレンダーでは、今話題のデータ関連SaaSを取り上げていきますので、引き続き乞うご期待ください！
