導入

DatabricksのVector Searchが2024/5月にGAとなりました。

その影響なのか、自分が利用している環境(AWS/東京リージョン)でも、Vector Searchエンドポイントが作れるようになっていました。いつの間に。

せっかくなので、以下の記事やLangChainのDocを読みつつ、LangChainと組み合わせた簡単なサンプルを実行してみます。

検証環境はDatabricks on AWS、DBRは15.1MLを使用しました。

Step1. パッケージインストール

必要なパッケージをインストールします。
重要なのはdatabricks-vectorsearchパッケージですね。

%pip install -U databricks-sdk
%pip install --upgrade --force-reinstall databricks-vectorsearch
%pip install -U langchain langchain_community

dbutils.library.restartPython()

Step2. サンプル文書を準備

ベクトル化する文書を準備するために、Wikipediaからテキストデータを取得・チャンク化し、LangChainのDocument形式に変換します。

import requests

def get_wikipedia_page(title: str):
    """
    Retrieve the full text content of a Wikipedia page.

    :param title: str - Title of the Wikipedia page.
    :return: str - Full text content of the page as raw string.
    """
    # Wikipedia API endpoint
    URL = "https://ja.wikipedia.org/w/api.php"

    # Parameters for the API request
    params = {
        "action": "query",
        "format": "json",
        "titles": title,
        "prop": "extracts",
        "explaintext": True,
    }

    # Custom User-Agent header to comply with Wikipedia's best practices
    headers = {"User-Agent": "tutorial/0.0.1"}

    response = requests.get(URL, params=params, headers=headers)
    data = response.json()

    # Extracting page content
    page = next(iter(data["query"]["pages"].values()))
    return page["extract"] if "extract" in page else None

full_document = get_wikipedia_page("葬送のフリーレン")

from typing import Any
from langchain.text_splitter import RecursiveCharacterTextSplitter

class JapaneseCharacterTextSplitter(RecursiveCharacterTextSplitter):
    """句読点も句切り文字に含めるようにするためのシンプルなスプリッタ"""

    def __init__(self, **kwargs: Any):
        separators = ["\n\n", "\n", "。", "、", " ", "", "==="]
        super().__init__(separators=separators, **kwargs)

# split it into chunks
text_splitter = JapaneseCharacterTextSplitter(chunk_size=400, chunk_overlap=80)
docs = text_splitter.create_documents([full_document], metadatas=[{"source":"葬送のフリーレン"}])
# print(docs)

Step3. Embedding用エンドポイントの準備

埋め込み計算に使うためのエンドポイントを準備します。
こちらは、以下の記事で作成したBGE M3モデルのサービングエンドポイントを利用します。

from langchain_community.embeddings import DatabricksEmbeddings

embedding_model_endpoint = "bge_m3_endpoint"

embeddings = DatabricksEmbeddings(endpoint=embedding_model_endpoint)

Step4. Databricks Vector Searchの準備

Vector Searchを利用するためのベクトル検索エンドポイント、およびインデックスを作成します。

まず、クライアントインスタンスを作成。

from databricks.vector_search.client import VectorSearchClient

vsc = VectorSearchClient()

ベクトル検索エンドポイントを作成。

vector_search_endpoint_name = "test-vector-search-demo-endpoint"

vsc.create_endpoint(
    name=vector_search_endpoint_name,
    endpoint_type="STANDARD"
)

数分後に利用可能となります。

次にインデックス(ベクトル変換後のデータなどを格納したオブジェクト)を作成します。

Databricksのインデックは以下の2種類あります。(こちらのドキュメントより)

種類	詳細
Delta同期インデックス	ソースDeltaテーブルと自動的に同期し、 Deltaテーブル内の基礎となるデータの変更に応じてインデックスを自動的に増分更新します。
ダイレクト・ベクトル・アクセス・インデックス	ベクターとメタデータの直接読み取りと書き込みをサポートします。ユーザーは、REST API または Python SDK を使用してこのテーブルを更新する必要があります。この種類のインデックスは、UI を使用して作成することはできません。 REST API または SDK を使用する必要があります。

実運用上はDeltaテーブルを基にインデックス作成することが多い＝Delta同期インデックスで作成することが多いと思います。
ただ、今回はChromaやFAISSなどと同様に、直接ベクトルストアを作成する「ダイレクト・ベクトル・アクセス・インデックス」で作成してみます。

VectorStoreClientのcreate_direct_access_indexを実行して、ダイレクト・ベクトル・アクセス・インデックスを作成します。

source_catalog = "training"
source_schema = "llm"
index_name = f"{source_catalog}.{source_schema}.demo_index"

index = vsc.create_direct_access_index(
    endpoint_name=vector_search_endpoint_name,
    index_name=index_name,
    primary_key="id",
    embedding_dimension=1024,
    embedding_vector_column="text_vector",
    schema={
        "id": "string",
        "text": "string",
        "text_vector": "array<float>",
        "source": "string",
    },
    embedding_model_endpoint_name=embedding_model_endpoint,
)

指定したテーブルにインデックスタイプ=直接アクセスのインデックスが作成されました。

Step5. インデックスへのデータ登録

ダイレクト・ベクトル・アクセス・インデックスは作成直後だと何もデータが入っていません。
Step2で作成した文書を登録しましょう。

LangChainのDatabricksVectorSearchクラスからインスタンスを作成します。

from langchain_community.vectorstores import DatabricksVectorSearch

dvs = DatabricksVectorSearch(
    index, text_column="text", embedding=embeddings, columns=["source"]
)

文書を追加。

dvs.add_documents(docs)

このあたり、LangChainはインターフェースが統一されているので、Chromaなどと同じ形で利用できますね。
これでインデックスにデータを追加できました。

カタログエクスプローラを見ると、登録されている件数を確認することができます。
（なお、サンプルデータ確認は未実装）

Step6. 検索

では、実際に検索してみます。
このあたりはLangChainの通常のVectorstore利用と同じやり方です。

query = "フェルンの声優は誰？"
dvs.similarity_search(query, k=1)

出力

[Document(page_content='「歴史上で最もダンジョンを攻略したパーティーの魔法使い」と自称するだけあり、ダンジョンには詳しい。道中で宝箱を発見するとその中身に異常なまでの興味を示し、判別魔法で99パーセントミミック（宝箱に化けた魔物）とみやぶってなお、残り1パーセントの可能性に賭けて宝箱を開け、上半身をミミックに噛まれてもがくという場面が何度も描かれている。\nフェルン (Fern)\n声 - 市ノ瀬加那', metadata={'source': '葬送のフリーレン', 'id': 'aa521aa0-5998-468a-9028-937945b1f1fc'})]

ちなみに一度作成したインデックスを再度利用するには、VectorSearchClientのget_indexからインデックスオブジェクトを再度取得できます。

index = vsc.get_index(endpoint_name=vector_search_endpoint_name, index_name=index_name)

# LangChainで利用する場合は、DatabricksVectorSearchで再度ラップ
dvs = DatabricksVectorSearch(
    index, text_column="text", embedding=embeddings, columns=["source"]
)

まとめ

いつ東京リージョンに来るかと心待ちにしていたのですが、さらっと来ていました。
これまでRAGを実装するにあたってChromaや外部のベクトルストア系サービスを使っていたものをDatabricksのサービスで置き換えることができそうです。

DatabricksのVector Searchで強力なのはUnity CatalogやDeltaテーブルの連携だと考えていて、従来のデータ管理の上に強力な検索コンポーネントが使えるというのはかなり良いです。
RAGはアーキテクチャ設計において、良い選択肢が増えてありがたいですね。