株式会社シベスピ 従業員ブログ

シベスピの社員ブログ。技術・想い・経験沢山書いていきます!

【Python】OpenaiのAPIとLlamaIndexで、自前のドキュメントから回答を生成してくれるプログラムを作ってみました

話題のChatGPT。
非常に優秀ですが、ネット上にない情報については正しい回答をしてくれません。
そこで、自分だけが持っているドキュメントを渡してその内容から回答を生成してくれるプログラムを作ってみました。

LlamaIndex

GPT(汎用技術)のようなLLM(Large Language Model:大規模言語モデル)にプライベートなデータを補強するために、
in-context learningという枠組みがあり、これを行うには
・データの取り込み
・インデックス化
が必要です。

そこで今回は、このデータの取り込み、インデックス化、
またそのインデックスを利用して質問に回答するところまでの
すべてを提供してくれるLlamaIndexを使用します。

ドキュメントを用意

今回は、ChatGPTに架空の用語を作ってもらって、
それをテキストファイルとして保存しました。

ソースコード

from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader
from llama_index import StorageContext, load_index_from_storage
from llama_index.node_parser import SimpleNodeParser
import os

# APIキーの設定
os.environ["OPENAI_API_KEY"] = '取得したAPIキー'

# 指定したディレクトリ内のドキュメントを読み込み
documents = SimpleDirectoryReader('data').load_data()

# 読み込んだドキュメントをノードに解析
parser = SimpleNodeParser()
nodes = parser.get_nodes_from_documents(documents)

# インデックスを作成・保存
index = GPTVectorStoreIndex(nodes)
index.storage_context.persist(persist_dir="index.json")

# インデックスのロード
storage_context = StorageContext.from_defaults(persist_dir="index.json")
index = load_index_from_storage(storage_context)

# 質問
query_engine = index.as_query_engine()
response = query_engine.query(input(">> "))
print(response)

実行例

ドキュメントに保存した架空の単語の説明を求めると、
記載内容を踏まえた回答をちゃんともらえました。

ドキュメント内に記載のない架空の単語についても聞いてみました。
わからないことはわからないと素直に回答するようです。

まとめ

ドキュメントを読み込ませてその内容を踏まえた回答を生成するプログラムを作成しました。

現状ではテキストファイルのみ読み取りを行っていますが、PDFやPowerPointなどにも対応させれば、より広い場面で使えるものになるかもしれないと思いました。