LoginSignup
10
7

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@chitomo12(Masatomo Fukuda)

Firestoreのベクトル検索で全文検索してみる

Posted at

目的

Firestoreの全文検索の実装にベクトル検索を使ってみて良かったので良さを共有したい

背景

趣味の個人開発でデータベースサイトを作りたいと思い立ち、そのためのデータベース基盤としてFirebaseのFirestoreを使うことにしました。

そこでテキストの全文検索機能を設計するにあたり、Firestoreの標準のクエリだと以下のような理由で厳しかったので何かしら別の策を考える必要が生じました。

  • Firestoreのクエリはドキュメント単位で課金されるため、ドキュメント総当たりでフィルタリングするとすれば課金額が膨大になる
  • Firestoreのクエリ演算子で検索に使えるものには完全一致（==）、配列内一致（array-contains）しかなく、検索キーワードとの類似度でソートするといったことができない
  • Elastic Search、Algoliaといった外部サービスを使えば一応全文検索できるが、これも結構コストがそこそこ掛かる

　
色々調べているうちに、Firestoreに最近追加されたベクトル検索を使えば行けそうだと分かったので、これを試してみることにしました。

ベクトル検索について

ベクトル検索とは

ベクトル検索とは、データを高次元のベクトル空間にマッピングし、類似性に基づいて情報を検索する技術です。
ZOZO TECH BLOGより

ざっくり噛み砕くと、テキストや画像などのデータを数字の集合（[-0.00949603, 0.021038500, -0.01994903, …]のような多次元行列）に変換し、集合同士の近似度を内積などで計算し、より近いデータを見つける手法です。

検索以外にもレコメンドシステムでも使われているようなので、応用性の高い技術だと言えます。

協調フィルタリングとベクトル検索エンジンを利用した商品推薦精度改善の試み | メルカリエンジニアリング

K最近傍（KNN）法

  • ベクトル検索のアルゴリズムの一手法です。詳しくは↓

    ChatGPTによる解説

    K近傍法（K-Nearest Neighbors, KNN）は、分類や回帰に使われるシンプルな機械学習アルゴリズムです。あるデータを分類する際、そのデータに最も近いK個の既知のデータ（近傍）のクラスを参照して、多数決で新しいデータのクラスを決定します。例えば、K=3の場合、最も近い3つのデータ点のクラスのうち多数を占めるクラスに新しいデータを分類します。

    KNNの特徴は、距離（例えばユークリッド距離など）を基にして近いデータを見つけるため、訓練時にモデルを作成するのではなく、予測時に計算を行う「遅延学習」の手法であることです。また、Kの値の選び方やデータのスケーリングが性能に影響を与えます。

ベクトル化について

  • テキストを意味や類似度で分類されたベクトルに変換するには機械学習モデルが必要です
  • クラウド上でPythonを動かせるFirebase Cloud Functionsに機械学習用のフレームワーク（torch、transformersなど）とモデルをデプロイしようとしたところ、Cloud Functionsのメモリ制限に引っ掛かってしまい断念
  • 下記ブログでOpenAI APIのベクトル埋め込み用のtext-embedding-3-smallモデルを使う方法が紹介されており、Functionsにフレームワークやモデルをデプロイしなくて良い（容量が節約できる）のと、コスト面で非常に安価[^2]だったので今回はこれを採用しました

いざ実装

全体フロー図

Firebase

1. ReactのプロジェクトにFirebaseとFirestoreを追加

公式ドキュメントなどを参照してセットアップしてください。
こちらの記事でも分かりやすくまとめられています。

2. Cloud Functionsをセットアップ

下記記事などを参考に、各環境に合わせてセットアップしてください。

Cloud Functionsは無料プランでは利用不可で、従量課金プラン（Blaze）に変更する必要があります。プランをアップグレードしても無料枠を越えなければお金はかかりませんが、コストを意識して運用するのも勉強になるので勇気を出してアップグレードしましょう。

3. Firestoreにベクトルインデックスを追加

Firestoreのコレクションに対してKNNベクトル検索のクエリを実行するにはコレクションのベクトルインデックスが作成されている必要があります。

任意のコレクションを作成して、以下の手順でベクトルインデックスを作成します。

ベクトルインデックス作成手順

  • （前提）google-cloud-sdkをインストールしていること

  • gcloud firestore indexes composite createコマンドでベクトルインデックスを作成（2,3分掛かる）
(base) user@MacBook-Air react-project % gcloud alpha firestore indexes composite create \
    --collection-group=インデックス対象コレクション名 \
    --query-scope=COLLECTION \
    --field-config field-path=embedding_field,vector-config='{"dimension":"1536", "flat": "{}"}' \
    --database="(default)"
  • --collection-group：インデックスを新規作成するFirestoreのコレクション名
  • --field-config
    • field-path：インデックス化したいベクトルデータが入ったフィールド名（今回はembedding_fieldとする）
    • vector-config
      • dimension：ベクトルの次元数。今回はOpenAIのtext-embedding-3-smallモデルを利用するので1536を指定（上限は2048）
      • flat：インデックスのタイプがflatであることを明記。他のオプションが無いのである意味おまじない。
    • database：対象とするデータベース。アプリにFirebaseプロジェクトが紐付けられていたら"(default)"でok
      　
  • 他のフィールドと掛け合わせた複合インデックスを作りたい場合は次のようなコマンドにする
gcloud alpha firestore indexes composite create \
    --collection-group=collection-group \
    --query-scope=COLLECTION \
    --field-config=order=ASCENDING,field-path="color" \
    --field-config field-path=vector-field,vector-config='{"dimension":"1024", "flat": "{}"}' \
    --database=database-id
  • 作成したインデックスはgcloud alpha firestore indexes composite list --database="(default)”で確認できる
(base) user@MacBook-Air react-project % gcloud alpha firestore indexes composite list --database="(default)"
---
fields:
- fieldPath: __name__
  order: ASCENDING
- fieldPath: embedding_field
  vectorConfig:
    dimension: 1536
    flat: {}
name: projects/react-study/databases/(default)/collectionGroups/sampleCollection/indexes/CICAgOi3kJAK
queryScope: COLLECTION
state: READY

OpenAI API

OpenAI APIキーを取得

こちらの記事を参考にAPIキーを取得してください。

Firebaseの従量課金プランと違い、こちらはプリペイドで料金をチャージし、使った分だけ残高から差し引かれる仕組みになってます。
大量にAPIリクエストをしない限り1円未満の請求で済む上、上限に達したところでAPIリクエストを受け付けなくなるだけです。うっかり事故になる心配が無いのは魅力的ですね。

APIリクエストとレスポンスは以下のような形式です。

request
curl https://api.openai.com/v1/embeddings \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $OPENAI_API_KEY" \
  -d '{
    "input": "Your text string goes here",
    "model": "text-embedding-3-small"
  }'
response
{
  "object": "list",
  "data": [
    {
      "object": "embedding",
      "index": 0,
      "embedding": [
        -0.006929283495992422,
        -0.005336422007530928,
        ...
        -4.547132266452536e-05,
        -0.024047505110502243
      ],
    }
  ],
  "model": "text-embedding-3-small",
  "usage": {
    "prompt_tokens": 5,
    "total_tokens": 5
  }
}

Cloud Functionsでのベクトル検索

ベクトル化したテキストをFirestoreに保存

  • Cloud Functionsで以下の処理を含む関数を定義します
    1. OpenAI APIのtext-embedding-3-smallモデルを使い、テキストをベクトル化
    2. Firestoreのコレクションにベクトル化したテキストを保存
const OPENAI_API_URL = "https://api.openai.com/v1/embeddings";
const OPENAI_API_KEY = "your_key_here";

// OpenAI APIコールとベクトル取得の共通関数
const getVectorEmbeddedText = async (text: string) => {
  const response = await axios.post(
    OPENAI_API_URL,
    { model: "text-embedding-3-small", input: text },
    {
      headers: {
        "Authorization": `Bearer ${OPENAI_API_KEY}`,
        "Content-Type": "application/json",
      },
    }
  );
  return response.data.data[0].embedding;
};

/* テキストをベクトル化するCloud Function関数（CORS設定は適当）*/
export const vectorEmbedText = onRequest(
  { cors: true }, 
  async (req, res) => {
  
      // POSTリクエストのボディからテキストを取得
      const { text } = req.body;
    
      try {
        // ベクトル化されたデータを取得
        const embedding = await getVectorEmbeddedText(text);
    
        // embeddingからFieldValue.vectorを生成
        const vector = FieldValue.vector(embedding);
    
        // コレクション 'sampleCollection' にドキュメントを追加
        const docRef = await db.collection("SampleCollection").add({
          text_field: text,
          embedding_field: vector,
          createdAt: admin.firestore.FieldValue.serverTimestamp(),
        });
        res.status(200).json({ success: true, id: docRef.id });
      } catch (error) {
        res.status(500).send("Error embedding text.");
      }
    }
);

上の関数を実行すると、Firestoreのコレクションにテキストとベクトルデータが追加されます。
成功すれば、下の画像のようにembedding_fieldにベクトルデータが登録されていることを確認できます。

スクリーンショット 2024-10-25 23.29.42.png

検索キーワードをベクトル化し、最近傍クエリを実行して結果を取得

  • こちらもCloud Functionsに次の関数を追加します
    1. OpenAI APIでテキストをベクトル化
    2. Firestoreのコレクションに対してベクトル検索クエリを実行
/* 最近傍クエリを実行する関数（CORS設定は適当）*/
export const getNearestNeighbor = onRequest({ cors: true }, async (req, res) => {
    try {
      // GETリクエストのクエリからテキストを取得
      const text = req.query.text;
      if (typeof text !== "string") {
        throw new Error("Invalid request: 'text' query parameter is required.");
      }

      // ベクトル化されたデータを取得
      const embedding = await getVectorEmbeddedText(text);

      // embeddingからFieldValue.vectorを生成
      const vector = FieldValue.vector(embedding);
      
      // ベクトル検索クエリを実行
      const vectorQuery: VectorQuery = db.collection("sampleCollection")
                                         .findNearest({
                                           limit: 5,
                                           distanceMeasure: "EUCLIDEAN",
                                           vectorField: "embedding_field",
                                           queryVector: vector,
                                       });
      const snapshot: VectorQuerySnapshot = await vectorQuery.get();

      // レスポンスをお好みに処理
      ...

実際の動作

書籍タイトルを何個か登録して実際に検索クエリの結果を見てみました。

vectorSearch.gif
　
登録データが少ないのでキーワードと無関係なドキュメントもヒットしますが、基本的には狙ったデータが上位に出るので非常に理想的な感じになりました🎉

Firestoreのドキュメント読み込み数もlimitで絞った分だけカウントされているのでそこまで気にしなくていい数値に収まっていました。

まとめ

メリット

  • 曖昧なキーワードでも期待通りのデータがヒットするのでとても良い✨
  • Firestoreの標準のクエリと同じようにシンプルに書けるので開発者経験的にもGood

デメリット

  • コレクション全体の中から最も近いものをフェッチするので、期待するデータがないと無関係な検索結果ばかり出てくる
    • 最近、Amazonとかでキーワードと関係無い商品がヒットするのもベクトル検索使ってるから？🤔
  • Firestoreの標準のクエリに備わったoffset、startAtなどクエリカーソルを操作するメソッドが使えないので、ページネーションを実現するにはlimitを増やして重複するデータを随時捌くしかない
    • （検索精度が高ければそもそもページネーション要らんのではと思わなくもない）

その他の学び

  • OpenAI APIのtext-embedding-3-smallモデルはお財布に優しい
    • 1000回検索して0.3円程度。コーヒー1杯で100万回検索できる☕

今後の課題

ベクトル検索は検索だけでなくレコメンドにも応用できるということだったので、今度はレコメンドエンジンの設計にも手を出してみたいです。

参考

10
7
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
10
7

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?