LlamaIndexを使ってローカル環境でRAGを実行する方法

こんにちは。電通総研コーポレート本部システム推進部の山下です。
最近はChatGPTなどのLarge Language Model(LLM)を利用したAIが話題ですね。
そのLLMを応用したRetrieval-Augmented Generation(RAG)という技術があります。
これは、LLMに文書検索などを用いて関連する情報を与えて回答させることで、
LLMが知識として持っていない内容を回答させたり誤った情報を答えてしまうハルシネーションを抑止する技術です。

今回はこのRAGをLlamaIndexというライブラリを使ってローカル環境で実装する方法について紹介します。

なぜローカル環境でLLMを利用したいのか

大変便利なツールのLLMですが、利用が難しいこともあります。
例えば、機密情報を取扱いたい、外部インターネットへの接続に制限が掛かっているといった場合です。
最終的にOpenAIなどの外部APIの利用が難しいという判断になる場合もあるかと思います。

こういった場面でローカル環境やオンプレミスのサーバ上でLLMを利用できればこういった問題を回避することが出来ます。

LlamaIndexとは？

今回はLLMを利用したアプリケーションを開発するためのフレームワークとしてLlamaIndexを利用します。LlamaIndexはどのようなライブラリなのでしょうか。
マニュアルの冒頭に以下のような説明があります。まさにRAGを実現するために作られたフレームワークです。

LlamaIndex is a data framework for LLM-based applications to ingest, structure, and access private or domain-specific data. It’s available in Python (these docs) and Typescript.

LlamaIndexは色々な機能を提供してくれますが、筆者が特に便利だと感じているのは
以下の機能です。

既存のデータの読み込みとインデックス化
各種LLMを利用したインタフェース(チャットインタフェースの実現)

これらの機能を組み合わせることで、既存のデータを読み込み、その情報を元にLLMへのクエリを作成してRAGの実現が可能となります。
その他のLlamaIndexの詳細な機能はマニュアルを参照してください。

環境構築

LLMを利用するにはGPUを利用した環境が望ましいです。
ここでは、Windows上のWSLとdevcontainerを利用した環境構築の方法について説明します。

また利用するGPUはNVIDIA社製のGPUを想定しています。
ここではRTX 3060を搭載している計算機で実行しました。

nvidia-dockerの環境構築

開発で利用するコンテナでGPUを利用可能にするため、 nvidia-container-toolkit のインストールを実施します。
ここでは WSL上のUbuntuで以下のコマンドを実行してインストールを実施します。

    distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
          && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
          && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
                sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
                sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
    sudo apt-get update
    sudo apt-get install -y nvidia-container-toolkit

Dev Containerの環境

次に、Dev Containerを利用してどこの環境でも手軽に開発が行えるようにします。
まず、ソースツリーのトップディレクトリに .devcontainer というディレクトリを作成します。今回は以下のようなファイル構成で作成します。

    .devcontainer
    ├── Dockerfile
    ├── devcontainer.json
    └── requirements.txt

まずは、 requirements.txt を以下のような内容で作成します。

torch==2.1.1
llama-index==0.9.13
transformers==4.35.2
llama_cpp_python==0.2.20

次に Dockerfile を以下のような内容で作成します。
llama_cpp_python がGPUを使うためにインストール時に=CMAKE_ARGS="-DLLAMA_CUBLAS=on"= を設定します。

FROM nvidia/cuda:12.3.0-devel-ubuntu22.04
RUN apt update \
        && apt install -y \
        wget \
        bzip2 \
        git \
        git-lfs \
        curl \
        unzip \
        file \
        xz-utils \
        sudo \
        python3 \
        python3-pip && \
        apt-get autoremove -y && \
        apt-get clean && \
        rm -rf /usr/local/src/*
    
# requirements.txt をコピーして必要なライブラリをインストール
COPY requirements.txt /tmp/
RUN pip install --no-cache-dir -U pip setuptools wheel \
   && CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install --no-cache-dir -r /tmp/requirements.txt

devcontainer.json の内容は以下のような内容で作成します。
runArgs に --gpus オプションを追加してコンテナ内部からGPUを利用可能にします。

{
  "name": "Cuda DevContainer",
  "build": {
    "dockerfile": "./Dockerfile"
  },
  "runArgs": ["--gpus", "all"]
}

ここまでで、Dev Containerを利用した開発環境の構築が出来ました。
それでは、さっそくLlamaIndexを利用したRAGの実装を行ってみましょう。
以降の作業は、構築したDev Container環境を起動して実施します。

LlamaIndexを利用したRAGの実装例

実装するシステムの概要

今回実装するRAGのシステムは、以下の仕組みのものを実装します。

data というディレクトリにあるテキストファイルを読み込む
テキストの内容を埋め込みモデルを利用してインデックスを作成する
インデックスとLLMモデルを利用して質疑応答を行う

また、今回利用するモデルは以下のものを利用します。

埋め込みモデル: multilingual-e5-large
LLMモデル: ELYZA-japanese-Llama

モデルのダウンロード

models というディレクトリを作成して、その中にモデルをダウンロードします。

git lfs clone  https://huggingface.co/mmnga/ELYZA-japanese-Llama-2-7b-instruct-gguf/ --include "ELYZA-japanese-Llama-2-7b-instruct-q8_0.gguf"

RAGの実装

ここまででRAGを実装する環境が整いました。
Pythonを使ってRAGを実装します。

必要なパッケージのimport

まずは llama_index などの必要なパッケージをimportします。

import logging
import os
import sys

from llama_index import (
    LLMPredictor,
    PromptTemplate,
    ServiceContext,
    SimpleDirectoryReader,
    VectorStoreIndex,
)
from llama_index.callbacks import CallbackManager, LlamaDebugHandler
from llama_index.embeddings import HuggingFaceEmbedding
from llama_index.llms import LlamaCPP

# ログレベルの設定
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)

data ディレクトリに含まれるテキストファイルを読み込みます。

# ドキュメントの読み込み
documents = SimpleDirectoryReader("data").load_data()

LLMのモデル、埋め込みモデルの初期化を行います。ここでは特に設定を変更することなく、multilingual-e5-largeをそのまま利用しています。ただし、公式サイトの説明に「query:」の接頭辞を付けて利用する方が良いと注意書きがあります。詳細に能力を確認する場合などは「query: 」を付けるような実装にする必要があります。

n_ctx はLLMモデルのコンテキスト数になります。ELYZAのcontext sizeは4096なので4096を設定しました。
n_gpu_layers は、LlamaCPPはCPUとGPUの両方を利用できるのですが、GPUでいくつの層を処理するのかを指定するものになります。RTX 3060では32程度を指定しておくのが良さそうだったので32を指定しています。

multilingual-e5-large は実行時にダウンロードされます。
このため初回実行時には時間がかかります。

ServiceContext を作成する時にしている、 chunk_size は分割した文書のサイズとなっています。この大きさはRAGを実施する際に与える文書の量に直接影響があります。
RAGではLLMに対する質問(クエリ)を構築する際に、コンテキスト情報として質問の内容に応じた文書を与える仕組みになっています。

この与えられる文書の大きさが chunk_size で決まります。1つの文書のサイズが大きい場合は前後の文脈を多く含むことになります。コンテキスト情報は多い方が良いので大きくしたくなります。しかし、1つのクエリに対して与えられる文書の量には上限があります。今回は4096が最大になっています(LLMに指定している context_size)。

また、検索でなどでコンテキスト情報を見つけて複数のチャンクを与えるということも可能です。
つまり、ヒットした文書のうち上位10件を与えるというような仕組みです。こちらもたくさん与えたくなります。しかしチャンクサイズと同様にコンテキストサイズには限りがあるためいくつチャンクを与えるのかは試行錯誤が必要となります。

今回は512を chunk_size として設定しました。これは、採用している Multilingual-E5-large モデルのおすすめのシーケンス長が512に
設定されているためです。

# LLMのセットアップ
model_path = f"models/ELYZA-japanese-Llama-2-7b-fast-instruct-gguf/ELYZA-japanese-Llama-2-7b-fast-instruct-q8_0.gguf"
llm = LlamaCPP(
    model_path=model_path,
    temperature=0.1,
    model_kwargs={"n_ctx": 4096, "n_gpu_layers": 32},
)
llm_predictor = LLMPredictor(llm=llm)
    
# 埋め込みモデルの初期化  
    # 埋め込みモデルの計算を実行するデバイスを指定。今回は埋め込みモデルをCPUで実施しないとVRAMに収まりきらないので、CPUで実施する。
# "cpu","cuda","mps"から指定する。
EMBEDDING_DEVICE = "cpu"

# 実行するモデルの指定とキャッシュフォルダの指定
embed_model_name = ("intfloat/multilingual-e5-large",)
cache_folder = "./sentence_transformers"

# 埋め込みモデルの作成
embed_model = HuggingFaceEmbedding(
    model_name="intfloat/multilingual-e5-large",
    cache_folder=cache_folder,
    device=EMBEDDING_DEVICE,
)

# ServiceContextのセットアップ
## debug用 Callback Managerのセットアップ
llama_debug = LlamaDebugHandler(print_trace_on_end=True)
callback_manager = CallbackManager([llama_debug])

service_context = ServiceContext.from_defaults(
    llm_predictor=llm_predictor,
    embed_model=embed_model,
    chunk_size=500,
    chunk_overlap=20,
    callback_manager=callback_manager,
)

# インデックスの生成
index = VectorStoreIndex.from_documents(
    documents,
    service_context=service_context,
)

ここまでで、RAGを実行するための準備は出来ました。

RAGで利用するプロンプトを設定します。

このプロンプトの中に {context_str} と書いてある部分が実行の際にはコンテキスト情報に
置き換えられます。 {query_srt} が質問内容で置き換えられます。

# 質問
temp = """
[INST]
<<SYS>>
以下の「コンテキスト情報」を元に「質問」に回答してください。
なお、コンテキスト情報に無い情報は回答に含めないでください。
また、コンテキスト情報から回答が導けない場合は「分かりません」と回答してください。
<</SYS>>
# コンテキスト情報
---------------------
{context_str}
---------------------

# 質問
{query_str}

[/INST]
"""

いよいよRAGを実行します。
query_engine を設定している箇所で、 similarity_top_k という引数があります。
これは幾つコンテキスト情報を与えるかという設定になります。ここでは5つのコンテキスト情報を与える設定にしてあります。

query_engine = index.as_query_engine(
    similarity_top_k=5, text_qa_template=PromptTemplate(temp)
)
   
while True:
    req_msg = input("\n## Question: ")
    if req_msg == "":
        continue
    res_msg = query_engine.query(req_msg)
    res_msg.source_nodes[0].text
    event_pairs = llama_debug.get_llm_inputs_outputs()
    print("\n## Answer: \n", str(res_msg).strip())

これで、RAGを用いて質疑応答ができるチャットが完成しました。
とても簡単に実装できましたね。

実行例

試しに実行してみましょう。今回は、太宰治の走れメロスを青空文庫からダウンロードしテキストファイルにして読ませて実行しました。

まずは1つ質問をしてみます。

## Question: メロスの友人の名前は？
...
(デバッグ用の情報が出力されます)
...
## Answer: 
セリヌンティウス

正解ですね。
さらにもう1つ質問をしてみます。

## Question: メロス メロスはどんな性格でしたか？
...
(デバッグ用の情報が出力されます)
...
## Answer: 
メロスは内気な性格だと読み取れます。
 
-   メロスは21歳の青年で、村の牧人であった。
-   メロスには父も、母も無い。女房も無い。十六の、内気な妹と二人暮しだ。
-   メロスは、村の或る律気な一牧人を、近々、花婿はなむことして迎える事になっていた。結婚式も間近かなのである。
-   メロスは、それゆえ、花嫁の衣裳やら祝宴の御馳走やらを買いに、はるばる市にやって来たのだ。

これらの事から、内気で控えめな性格だと読み取れます。

ちゃんとコンテキストを使って回答ができてそうですね。