はじめに
SIGNATEで開催された第３回金融データ活用チャレンジに、DataRobotを活用してRAGを構築し挑戦しました。DataRobotは数回のクリックでベクターDBやRAGを構築できてしまい、かつ様々なパラメータを簡単に調整することもできる優れものです。しかし、作成されたベクターDBって、どうチャンキングされているの？ と気にならないでしょうか...？
DataRobotは、DataRobotの様々な機能を操作可能なPythonクライアントを提供しています。例えば今回であれば、複数のパラメータのベクターDBを同時に作成するであったり、予測AIでも、複数のデータセットで一気にモデルを作成したいなどのユースケースにおいて、効率的に試行することができます。
そして、ベクターDBの中身を確認することのできるメソッドが備わっています！
実際にやってみる
前提として、DataRobotのコードスペースにノートブックを作成して作業をしていきます。またコードスペースを作成したワークベンチ内にベクターDBを構築している必要があります。
まずはライブラリをインストールしていきます。こちらのドキュメントで、"Upgrade pip to fix langchain installation issues" としていたので、pipとsetuptoolsのアップデートを一応しておきます。
!pip install --upgrade pip setuptools datarobot
インストールができたら、必要なライブラリをインポートしていきます。一つ目のdatarobot.models.genai.vector_database.VectorDatabaseが今回のメインなわけですが、詳細については以下のドキュメントから確認することができます。
from datarobot.models.genai.vector_database import VectorDatabase
import pandas as pd
以下を実行すると、作成したベクターDBの一覧を確認することができます。出力結果の各ベクターDBに割り振られているIDをこの後使います。
VectorDatabase.list()
# 以下のように結果が出力される
# [VectorDatabase(id=foo, name=bar, execution_status=COMPLETED, ...]
最後に以下を実行することで、ベクターDBによって作成されたチャンクと対応する埋め込みベクトルを含むparquetファイルをダウンロード、確認することができます。
# 確認したいベクターDBを取得
vdb = VectorDatabase.get("<ベクターDBのID>")
# 任意のパスにparquetファイルを保存
file_path = "chunks_and_embeddings.parquet.gzip"
vdb.download_text_and_embeddings_asset(file_path)
# pandasデータフレームで保存したparquetファイルを読み込み
df = pd.read_parquet(file_path)
dfを確認すると、以下のような結果を見ることができます。この結果を見ながら、チャンキングの戦略を考えていくことができます。
| # | text_chunks | source | page | start_index | embeddings |
|---|---|---|---|---|---|
| 0 | fogefoge | 1.pdf | 0 | 0 | [0.12345,0.12345...] |
| 1 | fugafuga | 1.pdf | 0 | 128 | [0.12345,0.12345...] |
| ... | ... | ... | ... | ... | ... |
アプリ化して使いやすくする
毎回コードを実行して確認するのも煩わしいので、アプリ化したくなりますよね？（強制）。せっかくなのでDataRobotに、みなさんお馴染みのstreamlitで作成したベクターDB確認アプリをホストしてみたいと思います。DataRobotへのカスタムアプリのホスティングは以下が参考になります。
まずは以下のアセットを用意します。
import streamlit as st
from datarobot.models.genai.vector_database import VectorDatabase
import pandas as pd
list_vdb = VectorDatabase.list()
selected_vdb = st.selectbox("ベクターDBを選択", list_vdb, format_func=lambda x: x.name)
vdb = VectorDatabase.get(selected_vdb.id)
file_path = "tmp.parquet.gzip"
vdb.download_text_and_embeddings_asset(file_path)
df = pd.read_parquet(file_path)
st.write(df)
#!/usr/bin/env bash
echo "Starting App"
streamlit run app.py
datarobot==3.6.2
streamlit==1.42.0
DataRobotのコードスペースでアプリ用のディレクトリを作成し、上記のアセットを作成したディレクトリ配下にアップロードします。以下はappディレクトリを作成した場合です。
DataRobotにアプリをホストする際に必要なCLIツール DRAppsをコードスペース内でインストールします。コードスペース内のターミナルで以下を実行します。
pip install git+https://github.com/datarobot/dr-apps
後はコマンドを実行するだけなのですが、コマンドを実行するためには、自身のAPIキーとホストする環境のIDを用意する必要があります。APIキーの取得については以下のドキュメントをご確認ください。
ホストする環境IDについては、[DataRobot] Python 3.12 Applications Baseという環境が、DataRobotから提供されているので、こちらを使用します。レジストリ > 環境 > [DataRobot] Python 3.12 Applications Base を選択すると、IDを取得できるのでコピーしておきます。
これでコマンド実行の準備が整いました。DataRobotのコードスペース内のターミナルで、アプリ用のディレクトリ（例だと~/storage/app）に移動して、コマンドを実行します。
cd ~/storage/app
drapps create -t <APIキー> -e "<環境ID>" -p . ベクターDB確認アプリ
処理が完了すると、レジストリ > アプリケーション から、ベクターDB確認アプリが無事ホストされていることが確認できます。
実際にアプリを開いてみるとこんな感じ。
まとめ
今回はDataRobotのPythonクライアントを使用して、ベクターDBの中身を確認する方法と、DataRobotにアプリをホストする方法についてご紹介しました。DataRobotはGUIでの簡単な操作に加えて、コードベースでより効率的な活用も期待できると感じました。DataRobot AI アクセラレータという、コードベースでDataRobotを活用した様々なレシピも公開されていましたので、色々触ってみると面白そうです