3分プロトタイピング: RAGを使ってAIチャットアプリケーションに知識を与える

連載「3分プロトタイピング」

  1. Streamlitを用いたAIチャットアプリ
  2. RAGを使ってAIチャットアプリケーションに知識を与える(この記事です)
  3. ベクトルデータベース超入門

前回の投稿でStreamlitを使ったAIとチャットするアプリケーションの雛形を作成しました。

tech.route06.co.jp

あのアプリケーションにくまモンについて聞いてみるとこんな回答が返ってきます。　

くまモンについて教えてください

それっぽいけど違いますね。今回は、このように特定のキャラクターや事象について正しい情報をAIに返してもらう方法を紹介します。説明が長くなるので3分を超えてしまいますが、コードは3分で書けるようになっていますので、早速やってみましょう。

AIに知識を教える3つの方法

まず、AIにキャラクターなどの「知識」教える3つの方法について紹介します。

  1. プロンプトエンジニアリング
  2. Retrieval Augmented Generation (RAG)
  3. Fine-tuning

プロンプトエンジニアリングは、AIに答えてほしい質問や要求を明確に伝えることで、正しい知識を教える方法です。例えば以下のように、くまモンの概要を教えて、これを使って回答するように指示すると、以降の会話でくまモンについての正しい情報を返してくれます。

プロンプトエンジニアリング

プロンプトエンジニアリングは、知識を教える最も簡単な方法ですが、プロンプトの許容量を超える知識を教えることはできません。また、知識を更新するためにはプロンプトを書き換える必要があり、アプリケーションとして運用していくのは難しいでしょう。

Retrieval Augmented Generationは、RAGと表記されることが多いです。これについはNVIDIAにわかりやすい説明があるので引用します。

Retrieval-Augmented Generation は、外部ソースから取得した情報を用いて、生成 AI モデルの精度と信頼性を向上させるテクノロジです。

生成 AI の最新の進歩を理解するために、法廷を想像してみてください。

裁判官は、一般的な法律の理解に基づいて審理し、判決を下します。時には、医療ミス訴訟や労働争議など、特定の専門知識が必要なケースもあるため、裁判官は裁判所書記官を法務図書館に送り、引用できる判例や具体的な事例を探させます。

優れた裁判官のように、大規模言語モデル (LLM) は人間の様々なクエリに答えることができます。しかし、出典を引用した信頼できる回答を提供するためには、モデルにも調査を行うアシスタントが必要です。

AI の裁判所書記官は、Retrieval-Augmented Generation 、略して RAG と呼ばれるプロセスです。 https://blogs.nvidia.co.jp/2023/11/17/what-is-retrieval-augmented-generation/

このページは、RAGの歴史や仕組みについて詳しく説明されていますので、興味がある方は読んでみてください。

最後のFine-tuningは、AIのモデルに「くまモン」を表現する大量のデータを与えて、AIにくまモンの特徴を学習させる方法です。これは、AIに知識を教える方法としては最も効果的ですが、データを用意するのが大変です。

今回はRAGを使って「くまモンについて教えてください」と聞いたら回答できるAIを作ってみましょう。

LlamaIndexでくまモンインデックスを作り、Streamlitと接続する

NVIDIAの説明に出てきた裁判所書記官は、ここではくまモンについて知っているくまモンインデックスです。まずはLlamaIndexを使ってこれを作り、前回作成したStreamlitアプリケーションと接続します。

LlamaIndexは、RAGのように独自データを使ってAIを使うユースケースを対象とした様々な機能を提供しているフレームワークです。

前回作成したStreamlitアプリケーションのディレクトリに移動して、LlamaIndexをインストールしてみましょう。

# 仮想環境を有効にする
source .venv/bin/activate
pip install llama-index setuptools

次にくまモンインデックスを作ります。私が作成したくまモンについて説明したテキストファイルがあるので、これをダウンロードして、LlamaIndexでIndex化します。

# テキストファイルをダウンロード
mkdir data
curl https://gist.githubusercontent.com/toyamarinyon/1b7da8c05ec30e0c1ce4d12616382de6/raw/d35c57509195629a39afaa3b2277107b9c38a99c/gistfile1.txt -o data/kumamon.txt
# Indexを作成するプログラムを実装する
touch build_index.py

プログラムは以下のコードをコピーして貼り付けてください。このコードはSnowflakeの方が公開しているブログ「Building a RAG based Blog AI assistant using Streamlit, OpenAI and LlamaIndex」を参考にしています。

import logging
import os
import sys
from shutil import rmtree

import openai
from llama_index import ServiceContext, SimpleDirectoryReader, TreeIndex
from llama_index.llms.openai import OpenAI

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

service_context = ServiceContext.from_defaults(llm=OpenAI())


def build_index(data_dir: str, knowledge_base_dir: str) -> None:
    """Build the vector index from the markdown files in the directory."""
    print("Building vector index...")
    documents = SimpleDirectoryReader(data_dir).load_data()

    index = TreeIndex.from_documents(documents, service_context=service_context)
    index.storage_context.persist(persist_dir=knowledge_base_dir)
    print("Done.")


def main() -> None:
    """Build the vector index from the markdown files in the directory."""
    base_dir = os.path.dirname(os.path.abspath(__file__))
    knowledge_base_dir = os.path.join(base_dir, ".kb")
    # Delete Storage Directory
    if os.path.exists(knowledge_base_dir):
        rmtree(knowledge_base_dir)
    data_dir = os.path.join(base_dir, "data")
    build_index(data_dir, knowledge_base_dir)


if __name__ == "__main__":
    main()

OpenAIのAPI Keyを環境変数に設定して実行してみましょう。

export OPENAI_API_KEY=YOUR_API_KEY
python build_index.py

以下のようなログが出力されれば成功です。

Building vector index...
Done.

Streamlitと接続する

作成したインデックスをStreamlitアプリケーションと接続します。前回作成したファイルに追記してもいいのですが、差分がわかりやすいように新しいファイルを作成します。

touch kumamon.py

以下のコードをコピーして貼り付けてください。

import streamlit as st
import os
import sys
import openai
from llama_index import StorageContext, load_index_from_storage

def load_index():
    """Load the index from the storage directory."""
    print("Loading index...")
    base_dir = os.path.dirname(os.path.abspath(__file__))
    dir_path = os.path.join(base_dir, ".kb")

    # rebuild storage context
    storage_context = StorageContext.from_defaults(persist_dir=dir_path)
    # load index
    index = load_index_from_storage(storage_context)
    query_engine = index.as_query_engine()
    print("Done.")
    return query_engine


st.title("Hello RAG!")

if "messages" not in st.session_state:
    system_prompt = (
        "Your purpose is to answer questions about specific documents only. "
        "Please answer the user's questions based on what you know about the document. "
        "If the question is outside scope of the document, please politely decline. "
        "If you don't know the answer, say `I don't know`. "
    )
    st.session_state['messages'] = [
        {"role": "system", "content": system_prompt},
        {"role": "assistant", "content": "何か気になることはありますか？"}
    ]

if "query_engine" not in st.session_state:
    st.session_state.query_engine = load_index()

for msg in st.session_state.messages:
    if msg["role"] not in ["user", "assistant"]:
        continue
    st.chat_message(msg["role"]).write(msg["content"])

if prompt := st.chat_input():
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)

    response = st.session_state.query_engine.query(prompt)
    st.session_state.messages.append({"role": "assistant", "content": f"{response}"})
    st.chat_message("assistant").write(f"{response}")

実行してみましょう。

streamlit run kumamon.py

preview

見た目は前回作成したものとほとんど同じですね。 では、くまモンについて聞いてみましょう。

くまモンについて教えてください

おお、正しそうです。もう答えてもらっていますが、作った人について改めて聞いてみましょう。

くまモンを作った人ついて教えてください

いい感じですね。

まとめ

今回は、RAGを使って、くまモンについての質問に答えてもらうアプリケーションを作りました。 これを応用すると、自分の会社のドキュメントでインデックスを作成すると、社内の質問に答えてもらうアプリケーションを作ることもできますね。ワクワクしますね。 ただ、その場合は、色々と工夫するところがあるので、また別の機会に紹介したいと思います。

謝辞

途中で引用したインデックスを作成するコードの他にも、私がこの投稿を書くにあたり、Snowflakeの方が書かれたブログ「Building a RAG based Blog AI assistant using Streamlit, OpenAI and LlamaIndex」を参考にさせてもらいました。ありがとうございます。