LangChain の QAGenerationChain によるドキュメントからのQA生成を試す

「LangChain」の「QAGenerationChain」によるドキュメントからのQA生成を試したので、まとめました。

1. QAGenerationChain

「QAGenerationChain」は、ドキュメントからQA生成するチェーンです。評価用のデータセットを作るのに役立ちます。

QA Generation — 🦜🔗 LangChain 0.0.141

2. ドキュメントの準備

はじめに、チャットボットに教える専門知識を記述したドキュメントを用意します。
今回は、マンガペディアの「ぼっち・ざ・ろっく！」のあらすじのドキュメントを用意しました。

・bocchi.txt

ぼっち・ざ・ろっく!（漫画） - マンガペディア

3. QA生成

Colabでの実行手順は、次のとおりです。

(1) パッケージのインストール。

# パッケージのインストール
!pip install langchain
!pip install openai

(2) 環境変数の準備。
以下のコードの <OpenAI_APIのトークン> にはOpenAI APIのトークンを指定します。(有料)

import os
os.environ["OPENAI_API_KEY"] = "<OpenAI_APIのトークン>"

(3) Colabにドキュメントを配置。
左端のフォルダアイコンでファイル一覧を表示し、ドキュメントをドラッグ&ドロップします。

(4) ドキュメントの読み込み。

from langchain.document_loaders import TextLoader

# ドキュメントの読み込み
loader = TextLoader("./bocchi.txt")
doc = loader.load()[0]

(5) チャンクに分割。

from langchain.text_splitter import CharacterTextSplitter

# チャンクに分割
text_splitter = CharacterTextSplitter(
    separator = "\n\n",  # セパレータ
    chunk_size = 1000,  # チャンクの最大トークン数
    chunk_overlap = 0,  # チャンクオーバーラップのトークン数
)
texts = text_splitter.split_text(doc.page_content)

(6) QA生成。
QAGenerationChain.from_llm()でQAを生成します。

from langchain.chat_models import ChatOpenAI
from langchain.chains import QAGenerationChain

# QAの生成
chain = QAGenerationChain.from_llm(
    llm=ChatOpenAI(temperature=1))
for text in texts:
    print(chain.run(text))

残念ながら英語の回答が混ざってました。

4. 日本語対応

プロンプトテンプレートをカスタマイズして、必ず日本語で回答するようにします。

(1) プロンプトテンプレートの作成。
元のプロンプトテンプレート(langchain/chains/qa_generation/prompt.py)のSystemに「answer in Japanese.」を追記したものを用意しました。

templ1 = """You are a smart assistant designed to help high school teachers come up with reading comprehension questions.
Given a piece of text, you must come up with a question and answer pair that can be used to test a student's reading comprehension abilities.
When coming up with this question/answer pair, you must respond in the following format:
```
{{
    "question": "$YOUR_QUESTION_HERE",
    "answer": "$THE_ANSWER_HERE"
}}
```

Everything between the ``` must be valid json.
answer in Japanese.
"""
templ2 = """Please come up with a question/answer pair, in the specified JSON format, for the following text:
----------------
{text}"""


from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

# プロンプトテンプレートの準備
prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessagePromptTemplate.from_template(templ1),
        HumanMessagePromptTemplate.from_template(templ2),
    ]
)

(2) プロンプト指定付きでQA生成。

from langchain.chat_models import ChatOpenAI
from langchain.chains import QAGenerationChain

# QAの生成
chain = QAGenerationChain.from_llm(
    llm=ChatOpenAI(temperature=1),
    prompt=prompt)
for text in texts:
    print(chain.run(text))

無事、日本語でQAが生成できました。

5. JSON出力

JSON出力するには、次のように記述します。

import json

# JSONに出力
items = []
for text in texts:
    item = chain.run(text)[0]
    items.append(item)
    print(item)
with open("qa.json", mode="w") as file:
    json.dump(items, file, ensure_ascii=False)

Unityではじめる機械学習・強化学習 Unity ML-Agents 実践ゲームプログラミング v2.2対応版