Zenn
🔖

RAG機能付きチャットボットを作ろう-6_VectorDBの設定

2024/12/07に公開
Python
Streamlit
LangChain
生成 AI
chromadb
tech

TL;DR

前回の記事では、チャット履歴をMarkdown形式で表示しました。本稿では

  • VectorDBの設定
  • PDFのベクトル化

を行います。

準備

インストール

VectorStoreとしてChromaDBを使います。以下でインストールします。

pip install chromadb
pip install langchain-chroma

PDFの格納フォルダ作成

文書情報はPDFで格納している方が多いと思いますので、本稿ではPDFの内容をベクトル化して、検索できるようにします。
まずは、PDFを格納するフォルダを作成します。
フォルダ名は自由ですが、本稿ではpdf_docsとします。
また、その中にサンプルのPDFファイルを格納します。ファイル名はgpt_review.pdfとします。
gpt_review.pdfにはChatGPTとは！？ChatGPTはOpenAIの開発した生成AIです。という文章を記載しています。

.
└── streamlit/
    ├── main.py
    └── .env
    └── pdf_docs/
        └── gpt_review.pdf

実装

主な変更点

主な変更点は以下です。

  • ベクトルDBの設定
  • PDFのベクトル化の関数の作成
  • ベクトルDBへの検索のためのretrieverの設定

コード

import streamlit as st
from openai import OpenAI
from dotenv import load_dotenv
import os
import chromadb
from langchain_chroma import Chroma
import openai
from pydantic import BaseModel
from typing import List
from langchain_openai import OpenAIEmbeddings
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import glob

# .envファイルから環境変数を読み込む
load_dotenv(".env")


# OpenAIのAPIクライアントを初期化

client = OpenAI(
    api_key=os.environ['OPENAI_API_KEY']
    )


# text splitterの定義
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=100,
    length_function=len,
    separators=["\n", " ", ".", ",", ";", ":", "(", ")", "[", "]", "{", "}", "<", ">", '"', "'", "、", "。", "，", "；", "：", "（", "）", "【", "】", "「", "」", "『", "』", "〈", "〉", "《", "》", "“", "”"],
    is_separator_regex=False,
)


# ベクトルDBのためのEmbeddingsの定義
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# Chromaの初期化
vector_store = Chroma(
    collection_name="example_collection",
    embedding_function=embeddings,
    persist_directory="./chroma_langchain_db",  # ベクトルDBの保存先
)


def pdf_to_vector():
    # 特定フォルダのpdfのリスト化
    path = "./pdf_docs"
    # pdf_list_stored.txt を読み込んで、pdf_list に格納
    pdf_list = []
    if os.path.exists("pdf_list_stored.txt"):
        with open("pdf_list_stored.txt", "r") as f:
            for line in f:
                pdf_list.append(line.strip())
    else:
        # pdf_list_stored.txt が存在しない場合は、ファイルを作成
        with open("pdf_list_stored.txt", "w") as f:
            pass


    # globを使ってファイル名のリストを取得
    # サブフォルダも含める場合は、"**/*.pdf"とする

    pdf_list_current = glob.glob(path + "/**.pdf")
    pdf_list_new = list(set(pdf_list_current) - set(pdf_list))


    docs = []
    for pdf in pdf_list_new:
        loader = PyMuPDFLoader(pdf)
        text = loader.load_and_split(text_splitter)
        docs.append(text)
        # Vector storeに保存
        vector_store.add_documents(text)
        pdf_list.append(pdf)

    # pdf_list を pdf_list_stored.txt に保存
    with open("pdf_list_stored.txt", "w") as f:
        for pdf in pdf_list:
            f.write(pdf + "\n")

## ベクトル検索のretireverの定義
retriever = vector_store.as_retriever(
    search_type="mmr", search_kwargs={"k": 1, "fetch_k": 10}
)
       

# プロンプトを入力すると、チャットボットが返答を返す関数を定義
# 入力はOpenAIのAPIクライアントとプロンプト
def default_chat(client, prompt):
    response = client.chat.completions.create(
        model="gpt-4o-mini", # 好きなモデルを選択
        messages=[
            {"role": "system", "content": "You are AI assistant."},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content

# streamlitのsession_stateにチャット履歴を保存する
# もしチャット履歴がなければ、空のリストを作成
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

# チャット履歴を表示する関数
def display_chat_history():
    for chat in st.session_state.chat_history:
        if chat["role"] == "user":
            st.markdown(
                # 背景をグレーにして、角を丸くする
                f'<div style="background-color: #f0f0f0; border-radius: 10px; padding: 10px;">'
                f"ユーザー: {chat['content']}"
                '</div>', unsafe_allow_html=True)

        else:
            st.markdown(
                # 背景を青にして、角を丸くする
                f'<div style="background-color: #cfe2ff; border-radius: 10px; padding: 10px;">'
                f"チャットボット: {chat['content']}"
                '</div>', unsafe_allow_html=True)

st.title('RAG機能付きチャットボットを作ろう')
st.write('streamlitを使ったUIの作成')

# チャット履歴を表示
display_chat_history()

prompt = st.text_area('プロンプト入力欄', )

button1, button2, button3 = st.columns(3)
if button1.button('チャット'):
    chat_response = default_chat(client, prompt)
    # チャット履歴に追加
    # ユーザーの入力を追加、roleはuser
    st.session_state.chat_history.append({"role": "user", "content": prompt})
    # チャットボットの返答を追加、roleはsystem
    st.session_state.chat_history.append({"role": "system", "content": chat_response})
    st.rerun()
if button2.button('RAG'):
    retriever_response = retriever.invoke(prompt)
    st.session_state.chat_history.append({"role": "user", "content": prompt})
    st.session_state.chat_history.append({"role": "system", 
                                          "content": retriever_response[0].page_content})
    st.write(retriever_response)
    st.rerun()

if button3.button('PDFをベクトル化'):
    pdf_to_vector()
    st.rerun()

かなり長くなりました。一つずつ確認していきましょう。

Text Splitterの定義

ここではlangchainのRecursiveCharacterTextSplitterを使って、テキストを分割します。
今回のPDFの中身は少ないので正直不要ですが、通常のPDFは複数ページにまたがることが多いので、分割しておくと便利です。

# text splitterの定義
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=100,
    length_function=len,
    separators=["\n", " ", ".", ",", ";", ":", "(", ")", "[", "]", "{", "}", "<", ">", '"', "'", "、", "。", "，", "；", "：", "（", "）", "【", "】", "「", "」", "『", "』", "〈", "〉", "《", "》", "“", "”"],
    is_separator_regex=False,
)

Embeddingsの定義

Embeddingとは文章をベクトル化することで、ここではOpenAIのEmbeddingsを使います。
複数のモデルが出ているので、好きなモデルを選択してください。

# ベクトルDBのためのEmbeddingsの定義
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

Chromaの初期化

すこしややこしいのですが、本稿ではchromadbとlangchainのChromaを使っています。
ここではChromaの方を使っています。
./chroma_langchain_dbのフォルダを作成して、ベクトルDBを保存します。
バージョンによっては、persist_directoryが別の表記になっているかもしれませんので、公式ドキュメントを参照してください。執筆時点で使用しているバージョンは langchain-Chroma 0.1.2 です。

# Chromaの初期化
vector_store = Chroma(
    collection_name="example_collection",
    embedding_function=embeddings,
    persist_directory="./chroma_langchain_db",
)

PDFのベクトル化

streamlitでは起動のたびにすべての処理が実行されるので、
PDFのベクトル化は「PDFをベクトル化」ボタンを押すと実行されるようにしています。
PDFの重複を避けるために、pdf_list_stored.txtに処理したPDFのリストを保存しています。

def pdf_to_vector():
    # 特定フォルダのpdfのリスト化
    path = "./pdf_docs"
    # pdf_list_stored.txt を読み込んで、pdf_list に格納
    pdf_list = []
    if os.path.exists("pdf_list_stored.txt"):
        with open("pdf_list_stored.txt", "r") as f:
            for line in f:
                pdf_list.append(line.strip())
    else:
        # pdf_list_stored.txt が存在しない場合は、ファイルを作成
        with open("pdf_list_stored.txt", "w") as f:
            pass


    # globを使ってファイル名のリストを取得
    # サブフォルダも含める場合は、"**/*.pdf"とする

    pdf_list_current = glob.glob(path + "/**.pdf")
    pdf_list_new = list(set(pdf_list_current) - set(pdf_list))


    docs = []
    for pdf in pdf_list_new:
        loader = PyMuPDFLoader(pdf)
        text = loader.load_and_split(text_splitter)
        docs.append(text)
        # Vector storeに保存
        vector_store.add_documents(text)
        pdf_list.append(pdf)

    # pdf_list を pdf_list_stored.txt に保存
    with open("pdf_list_stored.txt", "w") as f:
        for pdf in pdf_list:
            f.write(pdf + "\n")

この関数を以下のボタンを押すと実行されるようにしています。

if button3.button('PDFをベクトル化'):
    pdf_to_vector()
    st.rerun()

ベクトル検索のretireverの定義

ベクトルDBへの検索のためのretrieverを定義します。
kは返す文書の数、fetch_kは取得する文書の数です。
search_typeは検索の方法を指定します。ここではmmrを使っています。mmrを使うと、クエリに類似していて、かつ多様性のある文書を返します。
（この例ではpdfが一つなので意味はありませんが。）

retriever = vector_store.as_retriever(
    search_type="mmr", search_kwargs={"k": 1, "fetch_k": 10}
)

検索ボタンの処理

RAGの検索ボタンを押すと、retrieverを使って検索を行います。
なお、この例では類似文書の中身を取得しているだけなので、実際にはRAGにはなっていません。

if button2.button('RAG'):
    retriever_response = retriever.invoke(prompt)
    st.session_state.chat_history.append({"role": "user", "content": prompt})
    st.session_state.chat_history.append({"role": "system", 
                                          "content": retriever_response[0].page_content})
    st.write(retriever_response)
    st.rerun()

では、

streamlit run main.py

を実行して、動作を確認してみましょう。

まずはこのような画面が表示されます。PDFをベクトル化ボタンを押して、ベクトル化を行います。（先にPDFを格納しておいてください。）

次に、「ChatGPTとは？」と入力して、RAGボタンを押すと、以下のような画面が表示されたら成功です。

なお、チャットボタンを押したときとの違いを確認してみてください。

リンク

Discussion