🔍
Voyager と OpenAI Embeddings で類似文章の検索を行う
Spotifyの開発した距離計算ライブラリVoyagerを試したのでそのメモです
使用するライブラリの準備
pandas = CSVデータ読み込み・表形式データを扱う
openai = 文章埋め込み (text-embedding-ada-002)
import os
import pandas as pd
from openai import OpenAI
from voyager import Index, Space
os.environ["OPENAI_API_KEY"] = "sk-xxx"
検索対象の文章を埋め込む
今回は text-embedding-ada-002 を使用しました。埋め込むと1536次元になります。
client = OpenAI()
# OpenAI Embeddings を使用して text を埋め込む処理
def get_embedding(text, model="text-embedding-ada-002"):
text = text.replace("\n", " ")
return client.embeddings.create(input=[text], model=model).data[0].embedding
# CSVからデータの読み込みと埋め込み結果をDataFrameに追加する
df = pd.read_csv('data/articles.csv')
df_article = df[['id', 'body_text']]
df_article['embedding'] = df_article['body_text'].apply(get_embedding)
Voyagerインデックスを作成する
VoyagerのIndexは以下の3種類が選択でき目的の距離計算ロジックに合わせて選択する
- Space.Euclidean
- Space.InnerProduct
- Space.Cosine
# インデックスを作成し埋め込みベクトルを追加する
df_embedding = pd.DataFrame(df_headline["embedding"].tolist())
index = Index(Space.Euclidean, num_dimensions=df_embedding.shape[1])
index.add_items(df_embedding.values.tolist())
# 引数のテキストと近しい文章を返す(knn)処理
def get_similar_articles(text, top_k=3):
neighbors, distances = index.query(get_embedding(text), k=top_k)
print(neighbors, distances)
return df.iloc[neighbors]
df_result = get_similar_articles("東京特許許可局")
display(df_result[["id", "title", "body_text"]])
参考リンク
Discussion