LlamaIndexを使ってRAGを試す

モチベーション

Chatbot UIを試したこの投稿で、今後挑戦したいことの一つにRAG(Retrieval Augmented Generation)との連携と述べた。この投稿ではLlamaIndexを使ってRAGを実現する方法についてまとめた。

実は昨年末にLangchainを使ってRAGを試してみた。その後、LlamaIndexとのキーワードを良く聞いてきたので、今回LlamaIndexを使ってRAGを実現することにした。

情報源

  1. LlamaIndexを使ってローカル環境でRAGを実行する方法 この記事の内容(コード、RAGのための付加情報も含め)をそのまま使った。この記事は今年の1月に書かれたもので、LlamaIndexのバージョンアップする前のコードで記述されており、その部分は変更する必要あり。
  2. LangChainを使ったRAGをElyza 7bを用いて試したみた 昨年の年末休暇時に、RAGを試そうとして参考にした記事。
  3. ImportError: cannot import name ‘SimpleDirectoryReader’ from ‘llama_index’ (unknown location) llama-indexのバージョン変更に伴いライブラリの構成が見直されたため、過去のバージョンのpythonコードだとエラーが出る。
  4. Web Page Reader 情報源1.ではテキストデータを埋め込みモデルのデータとして使用したが、Webページの情報を埋め込みモデルにするヒントが書かれている。

LlamaIndexでの実装

Dockerfile

Dockerfileの前半部分は、情報源1.を参考にし、pythonライブラリを組み込む部分は、この投稿で使ったDockerfileから大部分を流用した。

「pip install llama-index」より前の2行は、llama_indexがv0.10.0以上の場合に必要となる。

FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

# Set bash as the default shell
ENV SHELL=/bin/bash

# Build with some basic utilities
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/*

# alias python='python3'
RUN ln -s /usr/bin/python3 /usr/bin/python

RUN pip install --upgrade pip setuptools \
	&& pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 \
	--index-url https://download.pytorch.org/whl/cu121 \
	&& pip install jupyterlab matplotlib pandas scikit-learn ipywidgets \
	&& pip install transformers accelerate sentencepiece einops \
	&& pip install langchain bitsandbytes protobuf \
	&& pip install auto-gptq optimum \
	&& pip install pypdf sentence-transformers \
	&& pip install llama-index-embeddings-huggingface llama-index-llms-openai \
	&& pip install llama-index-llms-llama-cpp \
	&& pip install llama-index

# Install llama-cpp-python[server] with cuBLAS on
RUN CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 \
        pip install llama-cpp-python[server] --force-reinstall --no-cache-dir

# Create a working directory
WORKDIR /workdir

# Port number in container side
EXPOSE 8888

ENTRYPOINT ["jupyter-lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token=''"]

CMD ["--notebook-dir=/workdir"]

上記のDockerfileで作成されたコンテナのpip listの抜粋。

llama_cpp_python                        0.2.75
llama-index                             0.10.38
llama-index-agent-openai                0.2.5
llama-index-cli                         0.1.12
llama-index-core                        0.10.38.post2
llama-index-embeddings-huggingface      0.2.0
llama-index-embeddings-openai           0.1.10
llama-index-indices-managed-llama-cloud 0.1.6
llama-index-legacy                      0.9.48
llama-index-llms-llama-cpp              0.1.3
llama-index-llms-openai                 0.1.20
llama-index-multi-modal-llms-openai     0.1.6
llama-index-program-openai              0.1.6
llama-index-question-gen-openai         0.1.3
llama-index-readers-file                0.1.22
llama-index-readers-llama-parse         0.1.4

torch                                   2.2.2+cu121
torchaudio                              2.2.2+cu121
torchvision                             0.17.2+cu121

pythonソースの変更箇所

Dockerfileのところでも触れたが、llama_indexがv0.10.0でライブラリ構成が変更になり、以前のコードだと次のとおりエラーになる。

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[1], line 5
      2 import os
      3 import sys
----> 5 from llama_index import (
      6     LLMPredictor,
      7     PromptTemplate,
      8     ServiceContext,
      9     SimpleDirectoryReader,
     10     VectorStoreIndex,
     11 )
     12 from llama_index.callbacks import CallbackManager, LlamaDebugHandler
     13 from llama_index.embeddings import HuggingFaceEmbedding

ImportError: cannot import name 'LLMPredictor' from 'llama_index' (unknown location)

そこで、次のように修正した。

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
"""
from llama_index.legacy import LLMPredictor
from llama_index.core.prompts import PromptTemplate
from llama_index.core import ServiceContext, SimpleDirectoryReader, VectorStoreIndex

from llama_index.core.callbacks import CallbackManager, LlamaDebugHandler
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.llama_cpp import LlamaCPP

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

LLMのパスを指定するf文字列の部分は、自分の環境に合わせて次のようにした。

model_path = f"../20240421_llamacpp/ELYZA-japanese-Llama-2-7b-fast-instruct-q4_K_M.gguf"

GPUを使うために、n_gpu_layersの箇所を

    model_kwargs={"n_ctx": 4096, "n_gpu_layers": -1},

さらにEMBEDDING_DEVICEの箇所を

EMBEDDING_DEVICE = "cuda"

に変更した。

その他、記事の説明では、chunk_sizeを512に設定したと、あったので以下の通り変更した。

    chunk_size=512,

その他は、情報源1.のコードをそのまま使った。

動作について

コードを実行すると、LLM(ggufで量子化されたElyzaモデル)を読み込み、RAGの埋め込みモデル、インデックスを作成し終わるまで、5分程度待った。

質問を投げると、2、3秒で回答が返ってきた。

回答を返す際には、GPU使用率は60〜70%程度だった。RTX A4000(16GB)、TITAN V(12GB)の両方で動作した。GPUメモリは9GB程度使用していた。

Webページから埋め込みモデルを作る

今回試した情報源1.では、テキストデータを埋め込みモデルに使う方法であった。WikipediaなどのWebページのデータを使う方法を情報源4.を参考に更に試した。

Webページとしては、r過程のページを使用した。次のセルを

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

以下のように変更した。

target_url = "https://ja.wikipedia.org/wiki/R%E9%81%8E%E7%A8%8B"
# NOTE: the html_to_text=True option requires html2text to be installed
documents = SimpleWebPageReader(html_to_text=True).load_data(
    [target_url]
)

実行結果は、次のとおり。

## Question:  中性子星合体におけるrプロセスのメカニズムを教えて
(中略)
## Answer: 
 中性子星合体の結果として、r過程が起こる可能性があります

中性子星は非常に高密度の状態で、その中で中性子捕獲が不安定な核が生成されます。この生成された中性子捕獲の不安定な核はベータ崩壊を繰り返しながら中性子をニッケル56のような核種に取り込むことで、rプロセスと呼ばれる過程を起こします

## Question:  銀河とブラックホールの共進化を説明して
(中略)
## Answer: 
 コンテキスト情報に無い情報は回答に含めません。

また、中性子星合体によるエネルギーの放出が最終的にブラックホールになるという情報はないので、この質問に対する回答は「分かりません」です。

rプロセスに関する質問については、少し専門的な回答が得られたが、銀河とブラックホールの共進化に関する質問のように、埋め込みモデルに無い内容・概念を質問すると、答えられなくなる。

今後

埋め込みモデルへの取り込みに関して、以下のような課題に挑戦していきたい。

  • 複数のWebページや特定のURL配下全体を埋め込みモデルに取り込む
  • PDFを埋め込みモデルに取り込む
  • テキスト、Web、PDFなど異なる形式のソースから埋め込みモデルに取り込む

llamaindexの構造について理解を深め、RAGを使った回答の精度を高めていきたい。