Zenn
Closed9

自前ドキュメントの Semantic検索を実装（Qdrant + OpenAI Embeddings）

OpenAI
Qdrant
yoshida567yoshida567

これは何？

Qdrantを使ったセマンティック検索機能を実装する記事。

想定ユースケース

  • サポートサイトでの記事検索フォームにセマンティック検索機能を実装する。
  • エンドユーザーが検索フォームに検索ワードを打ち込んだら、リアルタイムに10件程度の関連記事リンクと記事タイトルを表示したい。

先にまとめ

成果物は以下リポジトリにまとまっているので、結果だけ見たい場合は以下を参照すればOK。
https://github.com/Yoshida24/learn-qdrant/releases/tag/v0.2.0

このリポジトリには以下が含まれており、環境構築〜自前のデータを使ったセマンティック検索エンジンの作成まで一通り試せるようになっている。

  • Qdrant ローカル環境構築機能
  • セマンティック検索エンジン作成のサンプルデータセット
  • 自前で用意した記事をセマンティック検索エンジンへ流し込むためのバッチスクリプト
  • 簡単なセマンティック検索用インターフェース

関連記事

Qdrantのローカル環境構築 ... 以前書いた記事。ローカルにQdrantを使ったVectorDBを構築する手順を掲載した。本記事ではこの記事で触れているようなQdrantの環境構築が終わっていることを前提としている。ただし、上記成果物コードでは make setup で Qdrant の環境構築は自動的に終わるようになっているので、成果物コードを使う場合には Qdrantのローカル環境構築 は読まなくても大丈夫。

yoshida567yoshida567

この記事でやること

  • LangChain の Reference を検索できるようなセマンティック検索基盤を作成する。
  • LangChain の docs を全部インポートすると時間がかかりそうなので、今回は検索精度評価がしやすそうな langchain/docs/docs/use_cases 以下の記事だけに絞ってデータを投入し、これらの記事に対する検索ができればOKとする。
yoshida567yoshida567

実装の方針

以下の3つに分けて実装していく。

  • A. DBへのデータインポートを行うバッチの実装
  • B. 検索ワードを入力に取る検索機能の実装
  • C. 検索フォーム代わりのCLI

詳細は以下。

A. DBへのデータインポートを行うバッチの実装

B. 検索ワードを入力とした検索機能の実装

  • エンドポイントから入力ワードを与えたら処理を開始
  • 検索検索ワードを OpenAI Embeddings によって埋め込みベクトルに変換する
  • 埋め込みベクトル化された検索ワードをローカルの Qdrant に引き渡し、3件程度の類似記事検索を行う。
  • 検索記事のデータに含まれるメタデータから、以下の形式でデータを取得して返却する。
retval
{
    "data": {
        [
            "score": float,
            "article": { 
                "title": string,
                "url": string,
                "description": string
            }
        ]
    }
}

C. 検索フォーム代わりのCLI

検索フォームを実装するのは面倒なので、代わりにCLIを実装して手触り感を評価できるようにする。
なお、別にこれは必須で作らないといけないものではない。

入出力イメージは以下のようなものを想定している。

zsh
# ユーザーの入力
% 🔍 検索ワード: SQL 自動

# 検索結果の表示
🤖 Search results are:

1. Quickstart
- score: 0.1
- title: Quickstart
- url: https://python.langchain.com/docs/use_cases/sql/quickstart
- description: Quickstart In this guide we’ll go over the basic ways to create a Q&A chain and agent over a SQL database. These systems will allow us to...

# 以下同様に 3. まで続く
yoshida567yoshida567

実装を進めていく。

DBへのデータインポートを行うバッチの実装

検索したいドキュメントのソースデータを取得する。

今回検索に使うのは、LangChain Python版ドキュメント Use cases の記事。
2024/3現在では記事数は大体80記事程度。

まずはこのドキュメントのソースをローカルにダウンロードする。
Use cases のソースをローカルにダウンロード取得するには以下のようにする。

  • GitHub の LangChain のリポジトリを開く
  • ブラウザで . を押下して VSCode on GitHub を開く
  • docs/docs/use_cases を右クリックする
  • ディレクトリをダウンロードする

これでdocs/docs/use_casesのソースをダウンロードできる。ダウンロードしたファイル群のうちの大半のファイルはドキュメントのソースだが例外的に _category_.yml と言う名前のファイルだけは階層構造を表すためのファイルであるため不要である。そのため、この名前のファイルはDBに含める必要がないので削除しておく。ファイル数は3ファイル程度と少ないので手動で削除すればOK。

次に埋め込み用の前処理を行うが、今回はデータが ipynb と mdxしかなく前処理しなくても十分綺麗なので、前処理はスキップする。

次にOpenAI Embeddings によって、各記事のタイトル + URL + 本文 + ベクトル埋め込みがセットになったデータを生成する。コードは次。

embedding.py
import os
import glob
import pandas as pd
import re
from openai import OpenAI
import constants
from dataclasses import dataclass, asdict

from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from qdrant_client.http.models import PointStruct

from typing import Any


# ユーザー定義の変数
target_docs_dir = "src/semantic_search/docs_src/langchain/**"
root_in_target_docs_dir = "langchain"
qdrant_collection = "langchain-docs-semanticsearch-dev-0-2-0"


@dataclass
class Article:
    path: str
    file: str
    content: str


@dataclass
class DocsMetadata:
    title: str
    url: str
    content: str


@dataclass
class QdrantPayload:
    id: int
    vector: list[float]
    payload: dict[str, Any]


def load_files() -> dict[str, list[Article]]:
    """_summary_
    Example:

    ```python
    [
        {
            "path": "/docs/use_cases/apis",
            "file": apis,
            "content": "{\n "cells": [\n  {\n   "cell_type": "raw",\n ...",
        }
    ]
    ```

    Args:
        _no_args_
    Returns:
        dict[str, list[str]]: _description_
    """

    # Step 1: Get all files recursively
    files = glob.glob(target_docs_dir, recursive=True)

    # Step 2, 3 & 4: Extract path, file name and content
    articles = []
    for file in files:
        print(file)
        if os.path.isfile(file):
            parts = file.split("/")
            use_cases_index = parts.index(root_in_target_docs_dir)
            path_with_extention = "/" + "/".join(parts[use_cases_index + 1 :])
            path = re.sub(r"\.(.*)$", "", path_with_extention)
            file_name, _ = os.path.splitext(parts[-1])
            print(file_name)
            with open(
                file, "r"
            ) as f:  # Use the original 'file' variable with extension
                content = f.read()
            articles.append(Article(path, file_name, content))

    # Step 5: Sort and create dictionary
    articles = sorted(articles, key=lambda x: x.path)
    result = {"articles": articles}
    return result


# content -> embedding
def embedding(content: str) -> list[float]:
    api_key = os.getenv(constants.ENV_OPENAI_API_KEY)
    client = OpenAI(api_key=api_key)
    embedding = client.embeddings.create(
        model=constants.TEXT_EMBEDDING_MODEL,
        input=content,
    )
    return embedding.data[0].embedding


# embeddingしたデータを作成する
def create_qdrant_payload(articles: list[Article]):
    qdrant_payloads: list[QdrantPayload] = []
    for i, article in enumerate(articles):
        content_in_max_token = (
            article.content
            if len(article.content) <= constants.MAX_TOKEN_IN_TEXT_EMBEDDING_3_SMALL
            else article.content[: constants.MAX_TOKEN_IN_TEXT_EMBEDDING_3_SMALL]
        )
        embedded_content = embedding(content_in_max_token)
        url = f"https://python.langchain.com{article.path}"
        qdrant_payloads.append(
            QdrantPayload(
                id=i,
                vector=embedded_content,
                payload={
                    "title": article.file,
                    "url": url,
                    "content": article.content,
                },
            )
        )
        print(f"embedded {i}th article")
        import time

        time.sleep(1)
    return qdrant_payloads

create_qdrant_payload の戻り値が直接Qdrantに登録できるデータになっている。

yoshida567yoshida567

Qdrantへ埋め込みを登録する

上記スクリプトに下記を追加すれば埋め込みデータを登録できる。事前にQdrantをdocket経由で立ち上げておくことを忘れずに。

main.py
# Qdrantにembeddingしたデータを登録する
def register_embedding(qdrant_payloads: list[QdrantPayload]):
    # クライアントを初期化する
    client = QdrantClient("localhost", port=6333)

    # コレクションを作成する
    client.create_collection(
        collection_name=qdrant_collection,
        vectors_config=VectorParams(
            size=constants.DIMENSION_IN_TEXT_EMBEDDING_3_SMALL, distance=Distance.DOT
        ),
    )

    # ベクトルを追加する
    for qdrant_payload in qdrant_payloads:
        operation_info = client.upsert(
            collection_name=qdrant_collection,
            wait=True,
            points=[
                PointStruct(**asdict(qdrant_payload))
                for qdrant_payload in qdrant_payloads
            ],
        )

    return operation_info


def main():
    files = load_files()
    articles = files["articles"]
    print(pd.DataFrame(articles))

    qdrant_payloads = create_qdrant_payload(articles)
    print(pd.DataFrame(qdrant_payloads).head(5))
    print("embedding dimension: " + str(len(qdrant_payloads[0].vector)))

    operation_info = register_embedding(qdrant_payloads)
    print(operation_info)


main()
yoshida567yoshida567

検索機能の実装

検索ワードで記事検索をする機能を実装する。
大きな流れとしては以下。

  • 検索ワードを埋め込みベクトル化
  • Qdrantで類似検索
  • 類似度（スコア）の高いデータを複数返す

このときベクトルの次元が異なるとおそらく類似検索時に比較ができないので注意する。
実装した結果は次。

search.py
import os
from openai import OpenAI
import constants
from qdrant_client import QdrantClient

qdrant_collection = "langchain-docs-semanticsearch-dev-0-2-0"
search_limit_n = 2


# content -> embedding
def embedding(content: str) -> list[float]:
    api_key = os.getenv(constants.ENV_OPENAI_API_KEY)
    client = OpenAI(api_key=api_key)
    embedding = client.embeddings.create(
        model=constants.TEXT_EMBEDDING_MODEL,
        input=content,
    )
    return embedding.data[0].embedding


def search(query: str):
    # クライアントを初期化する
    client = QdrantClient("localhost", port=6333)

    # ベクトル化
    query_vector = embedding(query)

    # クエリを実行する
    # 近いベクトルをsearch_limit_nの数だけ検索する。
    search_result = client.search(
        collection_name=qdrant_collection,
        query_vector=query_vector,
        limit=search_limit_n,
    )

    return search_result

上記 main.py でデータの投入が終わっていれば、すでに search.py で検索ができる状態になっている。
簡単なCLIアプリを作って、検索できることを試してみる。

yoshida567yoshida567

簡単な検索アプリを作って使ってみる

CLIで検索クライアントアプリを作って検索を試してみる。
対話形式で検索ができるようなアプリを作ってみた。コードは以下。

repl_semantic_search.py
from search import search


def main():
    print("\n=== 🤖セマンティック検索を開始します。 ===\n")
    query = input(
        "検索ワードを入力してください。 (例:「チャットボット メモリ管理」「データ タグ 分類」「データベース 連携」)\n\n🔍 検索ワード: "
    )
    while True:
        response = search(query)
        print(
            f"""
あなたの検索ワード:
    - {query}
"""
        )

        print(
            f"""✅ 検索結果:
"""
        )
        for i, article in enumerate(response):
            print(
                f"""    - {i}th:
        - 記事タイトル: {article.payload["title"]}
        - 記事URL: {article.payload["url"]}
        - 類似度スコア(dot product): {article.score}
        - 記事ID: {article.id}
        - バージョン: {article.version}
"""
            )
        query = input(
            "\nお探しの結果は見つかりましたか？\n\nもう一度検索する場合、検索ワードを入力してください:\n\n🔍 検索ワード:"
        )


main()
yoshida567yoshida567

アプリで動作確認

以下のように検索できる。

% python src/semantic_search/repl_semantic_search.py

=== 🤖セマンティック検索を開始します。 ===

検索ワードを入力してください。 (例:「チャットボット メモリ管理」「データ タグ 分類」「データベース 連携」)

🔍 検索ワード: チャットボット メモリ管理

結果は次。最初に chatbot の memory management の記事が来ており、十分いい感じだと言える。
特に英語のドキュメントを日本語で検索できている点がすごい。

あなたの検索ワード:
    - チャットボット メモリ管理

✅ 検索結果:

    - 0th:
        - 記事タイトル: memory_management
        - 記事URL: https://python.langchain.com/docs/use_cases/chatbots/memory_management
        - 類似度スコア(dot product): 0.35029587
        - 記事ID: 2
        - バージョン: 74

    - 1th:
        - 記事タイトル: index
        - 記事URL: https://python.langchain.com/docs/use_cases/chatbots/index
        - 類似度スコア(dot product): 0.32134968
        - 記事ID: 1
        - バージョン: 74

ベクトル埋め込みの際に食わせたドキュメントが英語な上に .ipynb 形式だったため、日本語での検索を受け付けるかどうかに不明があったが、今回の検証ではベクトル埋め込みによる検索は言語によらず意味をとらえた検索ができるできるらしいことがわかった。これは多言語対応のあるサイトなどでは結構重宝しそう。
また、埋め込みに使ったモデルは text-embedding-3-small であって large ではないのだが、これでも十分いい感じにセマンティック検索ができていることもわかった。

yoshida567yoshida567

課題

  • このアプリケーションはETLのそれぞれの層が密結合になってしまっており、データセットの取得元変更が変わった際に都度書き換える必要性があるため再利用性が悪い。適切なモジュール単位に分けるべき。
  • Repository層を作るべき:Qdrantとの通信・Embeddings

展望

  • 埋め込みモデルは今回 text-embedding-3-small を使ったが、そのうち largeも試してみたい。
  • プロダクションで使うなら速度と品質は定量的な改善の取り組みが必要なので、以下のような指標を置いて、今後改善と検証を進めたい。

目標とする指標

  • 速度: ユーザーが検索フォームに検索ワードを打ち込んだら、keyup のタイミングから 100ms[1] 以内に10件程度の関連記事データが返却される。
  • 料金: 1日あたり1万件のリクエストがあると仮定して、1月あたり（=20営業日）のリクエスト20万件に対して利用料金が$100を超えない程度を目標とする
  • 精度: この記事では精度評価に深入りしなかった。主観で見てそれなりの精度で検索できていればOKとしたが、精度を担保するための評価手法は研究しておくと納得して使ってもらいやすいと思う。
脚注

  1. 100ms以内という閾値: 100ms以内という閾値の決め方は『人間がWebページに対してなんらかの入力を行なった際に、100ms以内に応答が返ってくると「サクサク動作している」と感じる』という説に基づいている。これには厳密な根拠があるわけではないのだが、ゲーム開発を行っている企業の方がセミナーでそのように発表されていたため一旦この仮定に基づいて目標値を100msに決めた。実際100ms以内に応答が返って来れば体感的にも十分早いと感じられるので、一定納得のいく話だと思う。 ↩︎

このスクラップは2024/03/16にクローズされました