概要

LlamaIndexを使うと非常に簡単にRAG(Retrieval-Augmented Generation)を使った検索システムを作ることができます。

今回はLLMにない情報（PDF）をベクトル化して検索できる方法を紹介します。

環境

python 3.11.8
streamlit 1.31.1
llama-index 0.10.14

実装

開発環境

Python自体もそうですが、LangchainやLlamaIndexはバージョン更新のたびに破壊的な変更が多くバージョンを固定しないと期待通りに動かないことが多いです。

先日の記事のように最初に開発環境を整えることを推奨します。少なくともvenvやvirtualenvのような仮想環境は必ず使いましょう。

christina04.hatenablog.com

準備

まず必要なパッケージを追加します。

$ poetry add langchain-openai streamlit llama-index llama-index-llms-langchain

コード

今回のコードの完成形はこちらです。

import streamlit as st
import tempfile
from pathlib import Path
from langchain.chat_models import ChatOpenAI
from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader

index = st.session_state.get("index")


def on_change_file():
    if "index" in st.session_state:
        st.session_state.pop("index")


st.title("ベクトル検索")

# PDFをアップロードする
pdf_file = st.file_uploader("PDFをアップロードしてください", type="pdf", on_change=on_change_file)

if pdf_file:
    with st.spinner(text="準備中..."):
        # ファイルを一時保存する
        with tempfile.NamedTemporaryFile(delete=False) as tmp:
            tmp.write(pdf_file.getbuffer())
            reader = SimpleDirectoryReader(input_files=[tmp.name])
            documents = reader.load_data()
            Settings.llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
            index = VectorStoreIndex.from_documents(documents=documents)

if index is not None:
    user_message = st.text_input(label="質問を入力してください")

    if user_message:
        with st.spinner(text="検索中..."):
            query_engine = index.as_query_engine()
            results = query_engine.query(user_message)
            st.write(results.response)

非常に短いですが、これでRAGを実現できています。LlamaIndexというフレームワークの恩恵を感じますね。

ポイント

PDFアップロード

# PDFをアップロードする
pdf_file = st.file_uploader("PDFをアップロードしてください", type="pdf", on_change=on_change_file)

ファイルをベクトル化

SimpleDirectoryReaderでloadします。
今回はPDFですが、SimpleDirectoryReaderは様々なファイル(.csv, .docx, .md, .pdf, .mp3, .jpeg, .png, etc...)に対応しています。

with tempfile.NamedTemporaryFile(delete=False) as tmp:
    tmp.write(pdf_file.getbuffer())
    reader = SimpleDirectoryReader(input_files=[tmp.name])
    documents = reader.load_data()
    Settings.llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
    index = VectorStoreIndex.from_documents(documents=documents)

次にVectorStoreIndexにドキュメントを入れてindexを作ります。
VectorStoreIndexはドキュメントをノードに分割します。そして、LLMが照会できるように、各ノードのテキストのEmbedding（テキストのセマンティクスを数値化＝ベクトル化）を作成します。

indexに対してクエリ

LlamaIndexの用語ではindexはDocumentオブジェクトで構成されるデータ構造で、LLMによるクエリを可能にするものです。
ユーザ入力を使ってクエリを投げます。
クエリはそれ自体がEmbeddingに変換され、次にVectorStoreIndexによって数学的演算が実行され、どれだけ類似しているかによってランク付けされます。

query_engine = index.as_query_engine()
results = query_engine.query(user_message)
st.write(results.response)

動作確認

PDFの用意

適当にWikipediaから記事を参照し、PDF化します。
gpt-3.5-turboは2021年までのデータでトレーニングされているので、最近作られたものなどが良いでしょう。

今回は2024年１月にサービスリリースされた、パルワールドの記事を使ってみます。

ツール>PDF形式でダウンロードで保存できます。

検証

$ poetry run streamlit run home.py

でサーバが立ち上がります。

http://localhost:8501/

にアクセスします。

アップロード画面になっているので、先程のPDFを入れます。

ベクトル化しています。

質問

「パルワールドの発売日は？」のようなgpt-3.5-turboにない情報で質問します。

ちゃんと答えが返ってきています。

Wikiの情報の通りです。

その他

サンプルコード

今回のサンプルコードはこちら

github.com

ベクトル検索は類似度によるので正しい情報が返るとは限らない

ベクトル化して類似する箇所の内容から答えを生成するので、必ずしも正しい答え・期待する答えを返すわけではありません。

まとめ

LlamaIndexを使うと非常に短いコードでRAGを実現することができました。

Carpe Diem

備忘録

LlamaIndexでPDFに対してベクトル検索を行う

概要

環境

実装