こんにちは、ChatGPTに自社のデータや、専門的な内容のテキストに基づいて回答を作成して欲しいという需要はかなりあるのではないかと思います。
そうした用途のために、LangchainやLlama-indexを使用した解説が多く公開されていますが、OpenAIのcookbookにはライブラリを使わずにEmbeddingsを使用したサーチとクエリを実装する方法が解説されています。個人的な経験として、ライブラリベースで実装をすると、日本語のテキスト分割が微妙だったり、LLMの回答が英語になってしまったりと、余計に事態が複雑化して、なんだかなぁ〜という結果になりがちです。
この記事では、主に以下のドキュメントを参考にして、ベクトルデータベースにデータを保存するなど変更を加えています。間違いや、もっとこうした方がいいよ、などコメントありましたら、ぜひお願い致します。
作ってみるもの
Wikipediaのスティービー・ワンダーのページに記載されている情報を利用して、以下のような対話によるQ&Aを行います。
ユーザー：スティービーの武道館公演の回数を教えてください。
AI： スティービー・ワンダーの武道館公演の回数は、以下の通りです。
1975年：1月29日、30日
1981年：3月31日、4月1日
1982年：4日、5日
1995年：2月21日、22日
1996年：18日、19日、20日
合計で、スティービー・ワンダーの武道館公演は7回です。
手順
さて、スティービーに関する質問に答えてもらうために、いくつかの手順を踏んで実装してきます。
- データを準備する
Wikipediaからスティービー・ワンダーの記事を取得 - Embeddingを作成する
OpenAIのAPIを使ってEmbeddingを作成
ベクターデータベースに保存する - Q＆Aを行う
ユーザーの問い合わせを受けてベクターデーベースを検索
OpenAIのAPIに投げるメッセージを作成して、問い合わせを行う
ユーザーが質問を入力してから、回答が得られるまでに以下の経路をたどります。
ユーザーが質問を送信する ->
質問がベクトルデータに変換される ->
ベクターデータベースから類似度の高い文書が検索される ->
文書と質問をセットにしてAPIに問い合わせを行う ->
APIから回答が返却される
データを準備する
それでは、始めましょう。まずは、Wikipeidaからデータを取得するためのライブラリをインストールします。tiktokenはトークン数の計測のために使用します。
!pip install mwclient
!pip install mwparserfromhell
!pip install tiktoken
Wikipediaのページを取得し、Pandasのデータフレームに変換します。
import mwclient
import mwparserfromhell
import pandas as pd
def fetch_wikipedia_sections(title: str) -> pd.DataFrame:
site = mwclient.Site("ja.wikipedia.org")
page = site.pages[title]
wikitext = page.text()
parsed_wikitext = mwparserfromhell.parse(wikitext)
section_data = []
for section in parsed_wikitext.get_sections():
headings = section.filter_headings()
if not headings:
continue
heading_text = headings[0].title.strip()
section_text = section.strip_code().strip()
section_data.append({"heading": heading_text, "content": section_text})
return pd.DataFrame(section_data)
title = "スティーヴィー・ワンダー"
df = fetch_wikipedia_sections(title)
df = df[~df["heading"].isin(["脚注", "注釈", "出典", "参考文献", "関連項目", "外部リンク"])]
display(df)
以下が、Wikipediaから取得してきたサンプルデータです。外部リンクや、関連項目などは、スティービー自身と関連が少ない項目は除外しています。
Embeddingを作成する前にトークン数を測ってみます。
import tiktoken
GPT_MODEL = "gpt-3.5-turbo"
def num_tokens(text: str, model: str = GPT_MODEL) -> int:
encoding = tiktoken.encoding_for_model(model)
return len(encoding.encode(text))
df['num_tokens'] = df['content'].apply(num_tokens)
display(df)
ここではこのまま進めますが、GPT3.5の最大トークン数が4097であることを考えると、「経歴」のセクションはひとつの塊（チャンク）としては大きすぎるかもしれません。また複数の文脈が混在しているかもしれません。Embeddigという手法でGPTへの問い合わせに自前のデータを持ち込む場合、テキストをどのような塊に分割するかということが、最終的な回答の精度に大きく影響します。以下は、cookbookより一説を翻訳して引用します。
テキストをセクションに分割するための完璧なレシピはありません。
いくつかのトレードオフがあります：
- 長いセクションは、多くの文脈を必要とする質問に適している場合があります。
- 長いセクションは、より多くのトピックが混在している可能性があるため、検索に不利になる場合があります。
- 短いセクションは、コスト削減に有効です。（トークンの数に比例する）
- 短いセクションは、より多くのセクションを検索することができ、適切な回答の作成に役立つ可能性があります。
- オーバーラッピング（セクションの一部を重複させる）することで、セクションの境界で回答がカットされるのを防ぐことができます。
引用：Question answering using embeddings-based search
Embeddings
それでは、OpenAIのAPIを使用して、Embeddingsを作成します。
!pip install openai
!pip install chromadb
import os
import openai
os.environ["OPENAI_API_KEY"] = "[YOUR_API_KEY_HERE]"
openai.api_key = os.getenv("OPENAI_API_KEY")
EMBEDDING_MODEL = "text-embedding-ada-002" # OpenAI's best embeddings as of Apr 2023
BATCH_SIZE = 1000 # you can submit up to 2048 embedding inputs per request
def create_embeddings(items):
embeddings = []
for batch_start in range(0, len(items), BATCH_SIZE):
batch_end = batch_start + BATCH_SIZE
batch = items[batch_start:batch_end]
print(f"Batch {batch_start} to {batch_end-1}")
response = openai.Embedding.create(model=EMBEDDING_MODEL, input=batch)
for i, be in enumerate(response["data"]):
assert i == be["index"] # double check embeddings are in same order as input
batch_embeddings = [e["embedding"] for e in response["data"]]
embeddings.extend(batch_embeddings)
df = pd.DataFrame({"text": items, "embedding": embeddings})
return df
items = df["content"].to_list()
df_embedding = create_embeddings(items)
display(df_embedding)
無事にベクトルが作成できたことが確認できました。
次に、これらのデータを保存先を作成する必要があります。ここでは、オープンソースのEmbeddingデータベースであるChromaに保存し、Chromaに問い合わせることで類似した文書の検索を行います。
import chromadb
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
from chromadb.config import Settings
def create_chroma_client():
persist_directory = 'chroma_persistence'
chroma_client = chromadb.Client(
Settings(
persist_directory=persist_directory,
chroma_db_impl="duckdb+parquet",
)
)
return chroma_client
def create_chroma_collection(chroma_client):
embedding_function = OpenAIEmbeddingFunction(api_key=os.environ.get('OPENAI_API_KEY'), model_name=EMBEDDING_MODEL)
collection = chroma_client.create_collection(name='stevie_collection', embedding_function=embedding_function)
return collection
chroma_client = create_chroma_client()
stevie_collection = create_chroma_collection(chroma_client)
stevie_collection.add(
ids = df_embedding.index.astype(str).tolist(),
documents = df_embedding['text'].tolist(),
embeddings = df_embedding['embedding'].tolist(),
)
chroma_client.persist()
Chromaに対して問い合わせを行ってみましょう。「スティービーは日本武道館で何回公演している？」という質問に対して、ベクトル間の距離が近い（意味が近いと仮定される）文書がデータベースから返却されます。
def query_collection(
query: str,
collection: chromadb.api.models.Collection.Collection,
max_results: int = 100)-> tuple[list[str], list[float]]:
results = collection.query(query_texts=query, n_results=max_results, include=['documents', 'distances'])
strings = results['documents'][0]
relatednesses = [1 - x for x in results['distances'][0]]
return strings, relatednesses
strings, relatednesses = query_collection(
collection=stevie_collection,
query="スティービーは日本武道館で何回公演している？",
max_results=3,
)
for string, relatedness in zip(strings, relatednesses):
print(f"{relatedness=:.3f}")
display(string)
結果は以下の通り。日本公演のセクションがrelatedness=0.733と質問との関連が最も高く、次に経歴のセクションがrelatedness=0.679となりました。
Q&Aに答えてもらう
それでは、ここまで見てきた仕組みを利用して、GPTに送信するメッセージ(プロンプト)を作成します。ユーザーの質問から類似する文書を特定して、文書と質問のセットを作成し、OpenAIのAPIに問い合わせを行います。
今回の例では、ユーザーの質問文「スティービーの武道館公演の回数を教えて」と、関連度が最も高いwikipediaの「日本公演」をセットにしたメッセージが作成されます。
def query_message(
query: str,
collection: chromadb.api.models.Collection.Collection,
model: str,
token_budget: int
) -> str:
strings, relatednesses = query_collection(query, collection, max_results=3)
introduction = '以下の記事を使って質問に答えてください。もし答えが見つからない場合、「データベースには答えがありませんでした。」 と返答してください。\n\n# 記事'
question = f"\n\n# 質問\n {query}"
message = introduction
for string in strings:
next_article = f'\n{string}\n"""'
if (
num_tokens(message + next_article + question, model=model)
> token_budget
):
break
else:
message += next_article
return message + question
def ask(
query: str,
collection = stevie_collection,
model: str = GPT_MODEL,
token_budget: int = 4096 - 500,
print_message: bool = False,
) -> str:
"""Answers a query using GPT and a dataframe of relevant texts and embeddings."""
message = query_message(query, collection, model=model, token_budget=token_budget)
if print_message:
print(message)
messages = [
{"role": "system", "content": "スティービー・ワンダーに関する質問に答えます。"},
{"role": "user", "content": message},
]
response = openai.ChatCompletion.create(
model=model,
messages=messages,
temperature=0
)
response_message = response["choices"][0]["message"]["content"]
return response_message
それでは、質問してみましょう。
ask("スティービーの武道館公演の回数を教えてください。")
スティービー・ワンダーの武道館公演の回数は、以下の通りです。\n\n- 1975年：1月29日、30日\n- 1981年：3月31日、4月1日\n- 1982年：4日、5日\n- 1995年：2月21日、22日\n- 1996年：18日、19日、20日\n\n合計で、スティービー・ワンダーの武道館公演は7回です。
Wikipediaの情報を参考して、答えを返してくれました。うまく動いていますね。
print_message=Trueとすると、APIに送信されたメッセージを確認することができます。
ask("スティービーの武道館公演の回数を教えてください。", print_message=True)
以下の記事を使って質問に答えてください。もし答えが見つからない場合、「データベースには答えがありませんでした。」 と返答してください。
# 記事
日本公演
1968年
タムラ・モータウン・フェスティバル・イン・東京 2月13日 渋谷公会堂
1975年
1月29日、30日 日本武道館、31日、2月3日 大阪厚生年金会館、2月1日 静岡駿府会館
1981年
3月31日,4月1日 日本武道館
1982年
10月28日 福岡国際センター、29日 愛知県体育館、31日、11月1日、2日 フェスティバルホール、4日、5日 日本武道館、6日 横浜文化体育館、8日 宮城県スポーツセンター、9日 郡山市総合体育館
1985年
10月23日、24日、25日 大阪城ホール、27日 福岡国際センター、29日、30日、31日 国立代々木競技場第一体育館、11月2日、3日 後楽園球場、5日 仙台市体育館、7日、8日 道立産業共進会場この折り、札幌護国神社で宮司と記念撮影したものが、札幌市手稲のぎんれい写真館と札幌護国神社に残っている。
1988年
4月12日 福岡国際センター、13日 広島サンプラザ、15日、16日、17日 大阪城ホール、19日、20日 名古屋レインボーホール、21日 静岡草薙体育館、23日、24日 横浜スタジアム、25日、26日、27日 日本武道館、29日 仙台市体育館
1990年
12月13日、15日、16日 大阪城ホール、19日 名古屋レインボーホール、23日、24日 東京ドーム
1995年
2月21日、22日 日本武道館、24日、25日、26日 横浜アリーナ、3月1日、2日 大阪城ホール、9月4日、5日 マリンメッセ福岡
1996年
9月15日 横浜アリーナ、18日、19日、20日 日本武道館、21日 横田基地、23日 マリンメッセ福岡、24日 大阪城ホール、25日 名古屋レインボーホール、27日 石川産業展示館、29日 岩手産業文化センター
1999年
8月27日、28日 東京国際フォーラムホールA
2003年〜2004年
12月27日、28日 さいたまスーパーアリーナ、30日 マリンメッセ福岡
1月4日 名古屋レインボーホール、6日、7日 大阪城ホール
2007年
2月17日、18日 さいたまスーパーアリーナ、20日 名古屋レインボーホール、24日 宮城県総合運動公園総合体育館、27日、28日 大阪市中央体育館
2010年 サマーソニック2010
8月7日 大阪・舞洲アリーナ、8月8日 千葉マリンスタジアム
"""
# 質問
スティービーの武道館公演の回数を教えてください。
追記
続編として、Function Callingから今回作成したEmbeddingsのデータベースに問い合わせを行う方法を書いています。
参考