Transformers, LangChain & Chromaによるローカルのテキストデータを参照したテキスト生成

前回はHugging FaceのTransgormersとLangChainを用いたテキスト生成を実装しました。今回はさらにChromaを用いて、ローカルのDB上のデータを参照して質問応答を行うテキスト生成を実装してみます。

ChromaはいわゆるベクトルDBの一種です。ベクトルDBは埋め込みベクトルのような高次元のベクトルデータを扱うのに適したDBです。ベクトルDBの実装には色々とありますが、今回はサーバを構築しなくても簡単に試せるChromaを使ってみました。

このベクトルDBと大規模言語モデル（LLM: Large Language Model）を用いたテキスト生成を組み合わせることで、ユーザからの質問に関連した文章をベクトルDBから取得して、その文章を基に回答を生成する、といったことが可能になります。全体構成は次の図のようなイメージです。

インターネット上に公開していない（あるいは社外秘などで外部には公開できない）ドメイン固有の文書があり、生成AIを使ってそれらの文書をベースに質問応答を行いたい場合は、このような方法が有効になると思います。

実装の流れはこんな感じです。

ベクトルDBの構築
学習済みモデルをロード
タスクやモデルなどを指定してTransformers Pipelineを構築
PipelineとPromptTemplateを指定してLangChainのLLMChainを構築
質問文を入力してベクトルDBから類似度の高い文章を検索
推論の実行（LLMChainに質問文と検索結果を与えて回答を生成する）

なお、今回のコードはこちらです。

github.com

ベクトルDBの構築

今回使用する文書はIPAが公開している「アジャイルソフトウェア開発宣言の読みとき方」です。 PDFファイルをページ単位で埋め込みベクトルに変換してChromaに格納します。

LangChainにはPDFファイルを読み込むPDFMinerLoaderがあるので、それを使用して読み込みます。 PDFファイルを読み込むLoaderは他にもありますが、日本語が文字化けしたり、ページがきちんと認識できなかったりしたので、いくつか試した結果、PDFMinerLoaderを使うことにしました。ちなみに、名前のとおり、裏ではPDFMiner（正確にはpdfminer.six）が動いています。

from langchain.document_loaders import PDFMinerLoader

file = '000065601.pdf'

loader = PDFMinerLoader(file)
text = loader.load()
pages = text[0].page_content.split('\x0c')

ファイルをページ単位に分割できたら、HuggingFaceのsentence_transformers埋め込みモデルを使って、テキストデータを埋め込みベクトルに変換し、Chromaに格納します。今回も大規模言語モデルはサイバーエージェント社のOpenCALM-1Bを使います。

from langchain.embeddings.huggingface import HuggingFaceEmbeddings

model_name = 'cyberagent/open-calm-1b'
embeddings = HuggingFaceEmbeddings(model_name=model_name)

モデルをロードしたら、テキストデータを入力して埋め込みベクトルに変換します。 Chromaのオリジナルのインターフェースをそのまま使用しても良いですが、LangChainにChromaのラッパーがあるので、それを使うことにします。 langchain.vectorstores.ChromaでChromaのクライアントを生成します。 Chromaの引数embedding_functionに先ほどロードした言語モデルを指定します。引数persist_directoryを指定すると、ストレージ上の指定した場所にデータが保存されます。なお、引数persist_directoryを指定しない場合はインメモリとなります。

from langchain.vectorstores import Chroma

vectordb = Chroma(embedding_function=embeddings, persist_directory='./db')

続いて、先ほど読み込んだ文書について、ページ単位でループして、add_textsでDBに登録していきます。このとき、引数textsに渡したテキストデータが埋め込みベクトルに変換されて登録されます。

for i, page in enumerate(pages):
  if page == '':
    continue
  vectordb.add_texts(texts=[page],
                     metadatas=[{'source': file}],
                     ids=[f'id{i+1}'])

学習済みモデルをロード

学習済みモデルもサイバーエージェント社のOpenCALM-1Bを使います。なお、埋め込みベクトルを生成した際に使用したモデルと同じものを使う必要があります。

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)

Pipelineを構築

続いてTransgormersのPipelineを構築します。テキスト生成の場合はtransformers.pipelineの引数taskにtext-generationを指定し、modelとtokenizerに先ほどロードした学習済みモデルとトークナイザーを指定します。 kwargsの値はOpenCALM-1BのUsageを参考にしています。

from transformers import pipeline

task='text-generation'

kwargs = {
    'max_new_tokens': 64,
    'do_sample': True,
    'temperature': 0.7,
    'top_p': 0.9,
    'repetition_penalty': 1.05,
    'pad_token_id': tokenizer.pad_token_id
}

pipe = pipeline(
    task=task,
    model=model,
    tokenizer=tokenizer,
    device=device_id,
    torch_dtype=torch.float16,
    **kwargs
)

LLMChainを構築

langchain.llms.HuggingFacePipelineに先ほど構築したPipelineを指定します。

from langchain.llms import HuggingFacePipeline

llm = HuggingFacePipeline(pipeline=pipe)

また、PromptTemplateで大規模言語モデルに指示を出すためのプロンプトのテンプレートを作成します。テンプレートにはパラメータを含めることができ、{query}のような形式で記述します。のちほどテキスト生成を実行する際、このパラメータに値を渡します。

from langchain.llms import HuggingFacePipeline
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

llm = HuggingFacePipeline(pipeline=pipe)

template = """
質問: {query}

回答:
"""
prompt = PromptTemplate(template=template, input_variables=['query'])

LLMChainで大規模言語モデルとプロンプトを統合します。

llm_chain = LLMChain(llm=llm, prompt=prompt)

質問文に関連する文章をベクトルDBから取得

質問文を入力して、ベクトルDB上で質問文と類似度の高いページを検索します。 similarity_searchの引数kで上位何件までを返すのかを指定できます。

query = 'アジャイルソフトウェア開発宣言で、プロセスやツールよりも重視していることは？'
docs = vectordb.similarity_search(query=query, k=5)

ちなみにここでdocsの中身を確認してみると、質問文に関連したページが含まれていることが分かります。

推論

ベクトルDBの検索結果と元の質問文をLLMChainに渡して、最終的な回答を生成します。

llm_chain.run(input_documents=docs, query=query)

出力結果はこんな感じになりました。

'アジャイルは「俊敏な」「型にはまらない」という2つの言葉で表現されますが、「柔軟性・流動性が重要だ。またその柔軟性を保証するための仕組みが重要である。」というのが答えです。「変化への対応は遅いかもしれないけれど,すぐに改善するんだ」。これがチームにとって大切なことです。”スピード”と簡単にいいますが'

この回答だと、ベクトルDBから取得した文章を活用しきれていないように見えますが、大規模言語モデルのパラメータ数や文書を分割してChromaに登録した際の粒度（今回はページ単位）も影響しているかと思います。モデルを変えてみたり、文書の分割の粒度や分割方法を変えてみたりするなど、工夫の余地はありそうです。

noriho137’s diary

機械学習, 時々Web