概要
チャットボットのプロンプトを生成する際にあらかじめ用意しておいたQ&Aやトピックから、質問に対応するものを類似文検索したいことがあります。LlamaIndexを使えばあらかじめインデックス構築方法などが用意されていますが、LlamaIndexはアップデートが早く、維持するのが面倒なのと、もっとシンプルな実装が欲しかったのでopenaiモジュールのみで実装してみました。ここでは、ベクトル化をOpenAIで行っているため、openaiモジュールを使っていますが、コサイン類似度による類似文検索だけならnumpyで可能です。
背景
ChatGPT APIを使った簡単なチャットボットを作ってみましたが、その際に仮想人格を構成するために、過去の会話例や知識データを作成し、LlamaIndexで類似文検索をしていました。しかし、LlamaIndexがアップデートされたときに、APIパラメータのレベルで互換性が無くなり、LlamaIndexをそれほど深く使っていたわけではないのでもっとシンプルな実装にすることにしました。
単純な類似文検索(例えば、ユーザーの質問と過去の会話例の質問を類似文検索)をして、プロンプトに埋め込む程度で、Embeddingとベクトル検索ができれば十分、というケースでは有効だと思います。
ここではopenaiモジュールに組み込まれているnumpyを使ったコサイン類似度による検索を行っています。OpenSearch(Elasticsearch)を利用して、k-NN検索も可能ですが、外部サービスを立てるほどの規模ではない、という場合に使うと楽ができるかもしれません。
実装
- Python 3.9.14
- openaiモジュール 0.27.4
インデックスの構築
会話データ
テキストファイルで会話データを用意します。フォーマットは読み込めれば何でもいいですが、ここでは質問、回答を1行ずつで組にして、それぞれは空行で区切りました。
自己紹介して
私の名前はチャットちゃん。Web広報を担当してるよ。いろいろなことをとにかく試して楽しんでるよ。でも、たまにちょっとやり過ぎて怒られることもあるけどね。
自己紹介して
職場での役割は広報チームで広報全般を担当していて、施設公開などでは一般参加者への案内などをするよ。
どんなことしたの？
朝から一般の人の見学があったんだ。オフィスだけじゃなくて実験室も案内したよ。私たちは見慣れた景色だけど、やっぱり一般の人は装置が珍しくて楽しんでくれたみたいだよ。私も頑張って説明したよ。みんな楽
しんでくれたみたいだから、私も元気をもらえたよ！
どんな仕事してるの？
そろそろ施設公開について考える時期だから昨年度までの実績を検討しながらどんなことをするか話し合ったよ。
仕事をしてて悩みはある？
できれば研究者のみんなにももっと積極的に関わって欲しいんだけど、みんな忙しいからなかなか難しいのかな…。やっぱり、こういう説明って本当に一線で頑張っている人から直接話した方が面白いじゃない？
こんな感じで、50組ぐらい作成しました。
インデックス構築の実装
ここでは、質問のみをOpenAI Embedding APIを用いてベクトル化します。ベクトル化した情報を追加したものをリストにしてJSONとして出力すればインデックスは完成です。
#!/usr/bin/env python3
# チャットボット用会話インデックス
import json
import openai
talks = []
buf = open("talks.txt").read().strip()
for chunk in buf.split("\n\n"):
# 空行で区切って質問と回答に分ける
# embeddingを求めてオブジェクトにしてリストに追加していく
question, answer = chunk.split("\n")
res = openai.Embedding.create(input=[question], model="text-embedding-ada-002")
embedding = res["data"][0]["embedding"]
talks.append({
"question": question,
"answer": answer,
"embedding": embedding,
})
# リストをJSONとして出力する
json.dump(talks, open("talks.json", "w"), ensure_ascii=False, separators=(",", ":"))
ここで出力された talks.json がインデックスファイルです。
ベクトル検索
クラスを実装する
上記で用意したインデックスから検索するクラスを作成します。コサイン類似度を利用して検索するだけです。コサイン類似度計算は openai.embeddings_utils.cosine_similarity() がありますが、 openai.embeddings_utils は読み込むときにmatplotlibなど依存ライブラリが多いので、 cosine_similarity() のソースをコピーして使っています。
#!/usr/bin/env python3
import json, time
import openai
from openai.datalib import numpy as np
class SimpleVectorIndex:
EMBEDDING_MODEL = "text-embedding-ada-002"
def __init__(self, fullpath):
self.index = json.load(open(fullpath))
def cosine_similarity(self, a, b):
# openai.embeddings_utils.cosine_similarity()
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
def query(self, query_str, top_k=5, cutoff=None):
# 検索する
res = openai.Embedding.create(input=[query_str], model=self.EMBEDDING_MODEL)
embedding = res["data"][0]["embedding"]
similarities = []
for talk in self.index:
similarity = self.cosine_similarity(embedding, talk["embedding"])
if cutoff and similarity < cutoff: continue
similarities.append({"similarity": similarity, "talk": talk})
# ソートしてtop_kを返す
similarities.sort(key=lambda o: o["similarity"], reverse=True)
return similarities[:top_k]
if __name__ == "__main__":
idx = SimpleVectorIndex("talks.json")
start = time.time()
records = idx.query("自己紹介して", top_k=5, cutoff=0.8)
elapsed = int((time.time() - start) * 1000)
for rec in records:
print(rec["similarity"], rec["talk"]["answer"])
print(f"Elapsed: {elapsed}ms")
実行結果
$ time ./query.py
1.0 私の名前はチャットちゃん。Web広報を担当してるよ。いろいろなことをとにかく試して楽しんでるよ。でも、たまにちょっとやり過ぎて怒られることもあるけどね。
1.0 職場での役割は広報チームで広報全般を担当していて、施設公開などでは一般参加者への案内などをするよ。
0.8434879233468375 最近はマカロンとほしいもにはまってて、探している幻のほしいもマカロンがあるんだ！お菓子作りも好きで、自分で作ったお菓子を周りの人に振る舞うのが楽しいんだよね。
0.8397487043648575 そんな中で、私ができることは、専門的な内容をわかりやすく伝え、一般の方々にも興味を持ってもらうことなんだ。
0.8151114562156933 もっちろん！何だか急に仲良くなった感じがして嬉しいな！
Elapsed: 554ms
このようにシンプルに実装することができます。私の環境(Xeon Gold 6254 @ 3.10GHz)で検索にかかる時間が500msぐらいでした。
この結果をプロンプトなどに組み込むことでチャットボットとの会話をカスタマイズすることができます。
まとめ
LlamaIndexを使わずにコサイン類似度を用いた類似文検索を実装してみました。ここでは、質問に対してベクトル検索しましたが、同様に回答も検索したり、知識として持たせるトピックなどに応用することもできます。
この方法で速度が不足する場合はOpenSearchによるk-NN検索などを検討してもいいかもしれません。
参考
- OpenAI Documentation - Embedding
- Amazon OpenSearch ServiceでのK最近傍(k-NN)検索