はじめに
を参考にEmbeddingを用いて、ベクトル検索を実装したいと思います。
GAFAでは、ベクトル検索が主流らしく、また推薦システムに興味がある身としてコンテンツベース推薦とも親和性が高く、LLMの台頭により推薦と検索がより近い関係になるのでは?!という勝手な仮説(推薦・検索の互いを考慮したプラットフォーム開発に興味があるので)のもとベクトル検索に手を出します。
実装
titles = [
"World Briefings",
"Nvidia Puts a Firewall on a Motherboard (PC World)",
"Olympic joy in Greek, Chinese press",
"U2 Can iPod with Pictures",
"The Dream Factory",
"Fierce clashes kill three US soldiers, injure 14 in Afghanistan",
"US-EU talks on Airbus subsidies fail",
"Aussie alive after capture in Iraq",
"Israeli Helicopter Fires Missile in Gaza -Witnesses",
"Ballmer: We need a \$100 PC"
]
df = pd.DataFrame(data=[{"title": t} for t in titles])
df.head()
適当なデータフレームを用意します。ちなみにここでの「title」は、openai-cookbookリポジトリにある「AG_news_samples.csv」から10レコード分とってきました。
次に、「title」をEmbeddingにするための関数を作ります。
openai.api_key = os.environ["OPENAI_API_KEY"]
def get_embedding(
text:str,
model:str = "text-embedding-ada-002"
) -> List[float]:
embedding = openai.Embedding.create(input=text, model=model)["data"][0]["embedding"]
return embedding
まず、openAIのAPIを使うために、api_keyを発行する必要があります。(https://platform.openai.com/account/api-keys)
そして、モデルはopenAI推奨の「text-embedding-ada-002」を使います。このモデルはこれまで開発されてきたモデルの性能を多くの指標で超えているという報告があります(https://platform.openai.com/docs/guides/embeddings/what-are-embeddings)。
例えば、
・一度にインプットできるトークンの数が2000から約8000になったり(1文字あたり1~3 トークン 多分。。)
・出力されるembeddingの次元数が1536と減ってきている(類似度計算が速くなるなどのメリット)
など
そして、料金が大幅に安くなり、1000トークンあたり0.0004$と学生にも優しいです。
この関数に各行の要素を入れていきます。
df["embedding"] = df["title"].apply(lambda x: get_embedding(x))
df["embedding"] = df["embedding"].apply(np.array)
データフレームにapply関数を適用させ"embedding"カラムに入れます。また、後にnumpyで類似度を計算するので、numpy.ndarrayに変換しましょう。
次に、ベクトル検索のための関数を作りましょう。
def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
def vector_search(
query: str,
embeddings: List[np.ndarray],
k: int = 3,
distance_metric: str = "cosine"
) -> List[int]:
query_embedding = np.array(get_embedding(query))
distances = []
for i, item_embedding in enumerate(embeddings):
if distance_metric == "cosine":
cosine_distance = cosine_similarity(query_embedding, item_embedding)
distances.append((i, cosine_distance))
distances = sorted(distances, key=lambda x: x[1], reverse=True)[:k]
top_k = [d[0] for d in distances]
return top_k
vector_search関数について、検索クエリをqueryとして受け取り、これのembeddingを取得します。そして、各「title」のembeddingとのコサイン類似度を計算し、コサイン類似度の降順に並び替え、トップk件を返す。という流れになります。
ちなみに、コサイン類似度とは、
\begin{align}
Cossim(\boldsymbol{x},\boldsymbol{y}) &= \frac{\sum_{i=1}^n x_i y_i}{\sqrt{\sum_{i=1}^n x_i^2} \sqrt{\sum_{i=1}^n y_i^2}}
\end{align}
と、２つのベクトルのなす角のコサインの値を求めるものです。-1から1まで値を取り、ここでは、2つのembeddingが意味的に似ているとき、コサイン類似度は1に近くなります。なので、クエリと各embeddingとのコサイン類似度を求めて、降順に並び替えることでqueryと意味的に似ている文章を返せることが期待できます。
そして、これらの関数を適用させます。
embeddings = df["embedding"].to_list()
near_index = vector_search(query="iPod with Pictures",embeddings=embeddings)
df.iloc[near_index]
データ数が少ないので、結果の分散が大きい(学習はしてないので検索ヒットの意味で)と思いますが"iPod with Pictures"というクエリに対して、titleが "U2 Can iPod with Pictures"のレコードがトップにきています。
クエリを変えてみましょう。
embeddings = df["embedding"].to_list()
near_index = vector_search(query="Chinese press",embeddings=embeddings)
df.iloc[near_index]
"Chinese press"というクエリに対して、titleが "Olympic joy in Greek, Chinese press"というレコードがトップにきました。2行目以降がほとんど関係なさそうなので、もう少しデータを増やしたいですね。。
おわりに
Embeddings APIを使ってベクトル検索の雰囲気を感じてみました。今回はコサイン類似度を使ってベクトルの距離を計算しました。が、データ数が莫大になると、厳密にベクトルの距離を測っていてはユーザー体験を損なうことになります。近似最近傍探索という技術を使い少しアバウトに距離を計算することでレスポンス速度を速めることも注目を集めているようです。
また、自社データなど特定のデータをコンテキストとしてファインチューニングしたり、そもそも今回実装したコードをGPT-4に書かせたりと、LLMを活用する機会を増やして道具にしていきたい。