見出し画像

「Google Colab」でローカルLLMを用いたWeb検索RAGデモ

株式会社Lightblueインターンの堀川です。

 HuggingFaceで公開している「lightblue/suzume-llama-3-8B-multilingual」という弊社のモデルをファインチューニングすることによってWebで検索した情報をもとにRAGを行うことができるようにしたモデル「lightblue/duck_search_calling_multilingual_AWQ」のデモを作成いたしましたのでデモ環境のGoogle Colabのノートブックを公開いたします。

▶︎今回作成したデモのノートブック
(T4における動作を確認しています。)

作成のモチベーション

 Google社が公開しているGeminiなどのWebに接続することができるLLMが存在している一方で、ChatGPTに代表される多くのLLMでは回答に含まれる時間的な情報が最後に学習したデータに依存するため最新の情報が求められるタスクについて答えられない場合や答えたとしても誤った情報を出力する場合があります。

 当然ローカルのモデルではWebに接続することはできませんがWeb検索APIとファインチューニングされたモデルを組み合わせることでローカルモデルを用いてWeb検索した情報をもとに答えてくれるチャットボットを作成することを目指しました。

学習に使用したデータセット

 LangChainというLLMを用いたアプリケーション開発を効率的に行うためのライブラリではGoogle検索を元に回答を生成することができるようにする機能が備わっています。この機能およびgpt4を用いて多言語で行った入出力の記録である6137件のデータについて学習を行いました。

 具体的にはまずユーザの質問に対して検索のキーワードを返し、検索結果をまとめた入力により最後に回答を生成することを考えました。

 まず、検索のキーワードを返す応答を学習させるにあたり、特殊な形式であるということを学習させるために[KEY][/KEY]というスペシャルトークンで囲うこととしました。

 また、回答において情報の引用元のURLがリンクされていることを期待する一方で、URLはランダムな文字列が連続することもありLLMの回答として正確に出力することには限界があると考えました。そのためURLの文字列を[CITATION_1]といった短い表記に書き換えチャットボットにおいてさらにURLに書き戻すことにより正確にURLがリンクされるようにしました。

⬇︎実際入力したデータ例

<|im_start|>user
"リングワールド"の作者は誰ですか?<|im_end|>
<|im_start|>assistant
[KEY] "リングワールド"の作者 [/KEY]<|im_end|>
<|im_start|>user

[CITATION_0]
TITLE:リングワールド (ハヤカワ文庫 SF 616) | ラリイ ニーヴン, 小隅 黎 | 通販
DESCRIPTION: 前作の発表後、作者のニーヴンは非常に多くの読者から「リングワールドの力学的不安定さ」を指摘されました。この工学的問題の解決策に取り組んだニーヴンは、当初の予定に ...


[CITATION_1]
TITLE:リングワールド (架空の天体) - Wikipedia
DESCRIPTION: リングワールド (Ringworld) は、ラリー・ニーヴンのSF小説シリーズ〈ノウンスペース〉に登場する架空の巨大な人工天体。以下の4つの長編小説に登場し、それらの主な ...


[CITATION_2]
TITLE:リングワールドふたたび (ハヤカワ文庫 SF ニ 1-2 ノウンスペース ...
DESCRIPTION: ハードSFの楽しさを象徴する作品「リングワールド」 その続編が生まれた経緯が冒頭で語られます。前作の発表後、作者のニーヴンは非常に多くの読者から「リングワールドの ...


((以下引用省略))
<|im_end|>
<|im_start|>assistant
"リングワールド"の作者はラリー・ニーヴンです。彼はアメリカ合衆国の小説家で、特にSF作品で知られています。"リングワールド"は彼の代表作で、ヒューゴー賞、ローカス賞、ネビュラ賞を受賞しました。彼の作品は大胆なアイデアのハードSFを得意とし、しばしば推理小説と冒険小説の要素も含まれています。[詳細はこちら]([CITATION_3])<|im_end|>

copy

モデルの学習

 モデルの学習にはHuggingFaceと互換性があり、高速なLLMファインチューニングが可能である軽量ライブラリ、unslothを利用しました。

 また、Google ColabのT4環境でデモを回すためにモデルの量子化を行いました。具体的にはAutoAWQを使用して4bit量子化を行いました。

モデルの推論

モデルの推論を行うにあたりAWQに対応しており高速に推論できるパッケージとしてvllmを使用しました。

 しかし、T4環境で実行するにあたりメモリ容量が問題となってしまったため、モデルの最大トークン数を6000に、最大出力トークン数を512として設定しました。(おそらくそれでもWarningは出るかと思います、)

デモの使用方法

 今回作成したデモではGradioを使用してGUIとして直感的に使用できるようにしてあります。

 まず「事前準備」配下のセルを全て実行してください。必要なライブラリのインストールやモデルのダウンロードを行います。その後、「Run Gradio interface」配下の二つのセルを実行していただくとGradioのインターフェースが起動してデモを使用することができます。


 上記の例ではうまく来年大阪で開催される大阪・関西万博の情報を元に答えることができており、参考となるサイトを示すことができています。 試しに元のモデルである「lightblue/suzume-llama-3-8B-multilingual」で同様の質問をしてみると以下のような応答になります。

INPUT:
次の万博はいつ開催されますか

OUTPUT:
申し訳ありませんが、私には2023年時点での情報がありません。最新の会議やイベントの開催日の具体的な情報については、その関連する組織または公式ウェブサイトを確認してください。そのような特別なイベントに関連するニュースや声明が出た場合は、プレスリリースやオフィシャルソーシャルメディアチャンネルでもご確認いただくと良いでしょう。

copy

 このことからうまく元のモデルでは持っていなかった情報に対してWebでの検索結果を参考に回答を出力できるようになっていることが確認できました。

 一方で本モデルを使用するにあたって問題点が明らかになりました。例えば、上の例では「次の万博はいつ開催されますか」という質問に対して元のモデルが最初に「申し訳ありませんが、私には2023年時点での情報がありません。」と回答しているように、LLMでは自然な応答をするために最新の情報に対しては答えられない旨の出力をするように学習されています。そのため、検索モデルに対応して学習したにも関わらず最新情報に対する回答の制御の影響が出てきてしまうケースがあります。こういった場合には初めにスペシャルトークンで囲まれた検索のキーワードが出力されるはずが、最新の情報に対しては答えられない旨の出力が出てきてしまい、想定している回答のフォーマットから外れるためエラーが出てしまいます。

 これに対する改善策として質問の次の応答を出力させる際に出力の元のスペシャルトークンである[KEY] を入力している状態で続きを出力させることを考えました。すなわち以下のような入力を与えました。

USER: 次の万博はいつ開催されますか
ASSISTANT: [KEY]

copy

 この結果エラーが出る可能性を大幅に減らすことができました。

 また、他の問題点としては現在の日付に対しての知識がないため、例えば「今日は何月何日ですか」という質問に対し、Web検索で「今日は何月何日」というキーワードで検索が行われ、それに対して今日は何月何日というキーワードを含むWebサイトが複数与えられ、それぞれに書いてある日付が違う場合に適切なものを選ぶことができません。例えば、今回のデモで試した場合は「今日は2024年5月22日で、水曜日です。」という回答が出てきてしまいました(2024/6/10に試しました)。

最後に

 今回このようにWeb検索を通したRAGのデモを作成いたしました。まだまだ改善点はありますが、十分に使用できるモデルが作成できていると思います。

ぜひお試しください!

この記事が気に入ったらサポートをしてみませんか?