ChatGPTが当たり前のように使われるようになってから、ベクトル検索もまた当たり前のように使われるようになって久しいですが、どんなDBを使ってベクトル検索を実現するかは悩ましい問題ですよね。

AutoGPTやBabyAGIなどで外部のベクトル検索DBを指定する際にはPineconeが設定できるようになっているので、Pineconeを使う人も多いと思うのですが、有料利用が月70ドルからの価格設定になっているため、ちょっとしたアプリで使うにはチョット高いなー、と思ってしまいます。

OSSであればChromaDBという選択肢もありますが、推奨されているIn-memoryモードで運用しようと考えるとかなりメモリを食うのでインスタンスのサイズをそこそこ上げる必要がありますし、純粋にベクトル検索しかできないので、データストアとしては別のDBも用意しておく必要があります。

SQLiteという選択肢

そこでライトに別のDBを用意するとなるとSQLiteが検討にあがります。

SQLiteは軽量かつ組み込み型のリレーショナルデータベースで、以下のような特徴を持っています。

一方で、SQLiteの特性としてデータの書き込み時にはデータベース全体にロックがかかるため、高頻度にデータの書き込みが発生するようなWebアプリケーションに採用するのには向いていません。ただ読み取りトランザクションは並行して行うことができるため、書き込みはそこそこで読み取り主体のデータベース利用であればコスパ良く運用ができます。

公式サイトにも日次で10万ヒット未満のサイトであればSQLiteで問題なく動作するとしており、実証ではその10倍のトラフィックでも動作しているとしています。

大体のサービスではPostgreSQLなどといったクライアント／サーバ型のデータベースは、むしろオーバースペックなのかも知れません。

最近ではCloudflareがD1というSQLiteベースのマネージド分散データベースを発表していたり、Fly.ioでもLiteFSという、これもまたSQLiteを利用した分散データベースを開発しており、SQLiteにまた注目が集まっているように感じます。

SQLiteでベクトル検索を可能にするsqlite-vss

そんなポータブルで便利なSQLiteですが、そのSQLiteでベクトル検索ができるとなるとより夢が広がります。

SQLite自体はファイルベースなので、あらかじめベクトルデータを設定したSQLiteデータベースファイルをアプリに組み込んで配布しても良いわけです。そうすればデータベースサーバを用意しなくて済む分コストも圧縮されますし、組み込みなのでアプリからは軽量に動作します。

ホスティングする場合でもFly.ioのようにボリュームイメージを利用できるPaaSを利用すれば、問題なく運用が可能です。

前置きが長くなりましたが、このような夢を叶えてくれる拡張がsqlite-vssです。ベクトル検索はFaissベースで実装されています。

とっても良さげではあるのですが、実際に組み込んでみた場合のコード例が見つからなかったので、手を動かして試してみました。

実際にPythonから利用してみる

Jupyter Notebookから実行している前提でコードを掲載していきます。検証に利用しているPythonのバージョンは3.10.11です。

まずは必要なライブラリをインストールします。

!pip install -U numpy openai sqlite-vss

インストールした後、SQLiteのデータベースを作成します。拡張機能を有効にし、sqlite-vssがロードされるようにします。

import sqlite3
import sqlite_vss

db = sqlite3.connect('papers.db', timeout=10)
db.enable_load_extension(True)
sqlite_vss.load(db)
vss_version = db.execute('select vss_version()').fetchone()[0]
print('SQLite VSS Version: %s' % vss_version)

次にテーブルを作成します。今回は論文のアブストラクトや要約データを登録し、要約データをベクトル検索できるようにする、というシナリオでデータを作成していきます。

# papersテーブルの作成
db.execute('''
    CREATE TABLE IF NOT EXISTS papers(
        id INTEGER PRIMARY KEY,
        title TEXT,
        abstract TEXT,
        summary TEXT,
        url TEXT,
        created_at DATETIME
    );
''')

埋め込みベクトルを保存するテーブルは仮想テーブルで用意します。この仮想テーブルには通常のカラムを設定できないので、ROWIDにpapersテーブルのidを設定する形でリレーションを張る想定で進めます。

# vss_papersテーブルの作成
db.execute('''
    CREATE VIRTUAL TABLE IF NOT EXISTS vss_papers USING vss0(
        summary_embedding(1536)
    );
''')

次にデータを登録するための関数を用意します。

今回埋め込みベクトルはOpenAIのtext-embedding-ada-002を利用して作成していますが、このAPIで返ってくるList[float]型のデータを直接埋め込みベクトルとして保存することはできないため、numpyを利用してraw bytesフォーマット（bytes型）に変換しています。

from typing import List
import numpy as np
import openai
from datetime import datetime


def serialize(vector: List[float]) -> bytes:
  """ serializes a list of floats into a compact "raw bytes" format """
  return np.asarray(vector).astype(np.float32).tobytes()

def generate_embedding(text: str) -> List[float]:
    response = openai.Embedding.create(
        engine="text-embedding-ada-002",
        input=[text]
    )
    return response['data'][0]['embedding']

def insert_paper(title, abstract, summary, url):
    current_time = datetime.now()
    summary_embedding = generate_embedding(summary)

    with db:
        db.execute('''
            INSERT INTO papers(title, abstract, summary, url, created_at)
            VALUES (?, ?, ?, ?, ?)
        ''', (title, abstract, summary, url, current_time))

        last_id = db.execute('SELECT last_insert_rowid()').fetchone()[0]

        db.execute('''
            INSERT INTO vss_papers(rowid, summary_embedding)
            VALUES (?, ?)
        ''', (last_id, serialize(summary_embedding)))

papersテーブルにレコード追加するのと同時に、埋め込みベクトルを保存しているvss_papersテーブルにもレコード追加を行っています。

更にこれらのテーブルをジョインしてベクトル検索できる関数を用意すれば簡易的に利用するための準備が完了です。引数「k」は何レコード抽出するかを決定するパラメータです。

def search_similar_embeddings(query_embedding, k=5):
    results = db.execute('''
        SELECT papers.*, vss_papers.distance
        FROM vss_papers
        JOIN papers ON vss_papers.rowid = papers.id
        WHERE vss_search(vss_papers.summary_embedding, vss_search_params(?, 10))
        ORDER BY vss_papers.distance
        LIMIT ?
    ''', (serialize(query_embedding), k))
    return results.fetchall()

動作確認のためにテストデータが必要なので、それっぽいデータをChatGPTに作成してもらいました。

papers = [
    {
        'title': '深層学習による画像認識の進化',
        'abstract': '深層学習技術を用いた画像認識の最新の進展について述べる。',
        'summary': '深層学習は画像認識の精度を大幅に向上させ、AIの新たな可能性を開いている。',
        'url': 'http://example.com/paper1',
        'created_at': '2023-06-01 00:00:00'
    },
    {
        'title': '量子コンピューティングの可能性',
        'abstract': '量子コンピューティングの基本的な概念とその可能性について解説する。',
        'summary': '量子コンピューティングは、現在のコンピューティング能力を大幅に超える可能性を秘めている。',
        'url': 'http://example.com/paper2',
        'created_at': '2023-06-02 00:00:00'
    },
    {
        'title': '自然言語処理の最新動向',
        'abstract': '自然言語処理技術の最新の動向とその応用について概説する。',
        'summary': '自然言語処理は、人間とコンピュータとのコミュニケーションをより自然にするための重要な技術である。',
        'url': 'http://example.com/paper3',
        'created_at': '2023-06-03 00:00:00'
    },
    {
        'title': 'ブロックチェーン技術の進化',
        'abstract': 'ブロックチェーン技術の進化とその社会への影響について考察する。',
        'summary': 'ブロックチェーンは、透明性とセキュリティを確保するための革新的な技術である。',
        'url': 'http://example.com/paper4',
        'created_at': '2023-06-04 00:00:00'
    },
    {
        'title': '人工知能の倫理問題',
        'abstract': '人工知能の発展に伴う倫理的な問題とその対策について議論する。',
        'summary': '人工知能の発展は、新たな倫理的な問題を引き起こす可能性がある。',
        'url': 'http://example.com/paper5',
        'created_at': '2023-06-05 00:00:00'
    },
    {
        'title': 'データサイエンスの未来',
        'abstract': 'データサイエンスの未来の展望とその社会への影響について述べる。',
        'summary': 'データサイエンスは、情報化社会における意思決定を支える重要な分野である。',
        'url': 'http://example.com/paper6',
        'created_at': '2023-06-06 00:00:00'
    },
    {
        'title': '仮想現実と拡張現実の可能性',
        'abstract': '仮想現実（VR）と拡張現実（AR）の可能性とその応用について考察する。',
        'summary': 'VRとARは、新たな体験と効率的な作業を可能にする革新的な技術である。',
        'url': 'http://example.com/paper7',
        'created_at': '2023-06-07 00:00:00'
    },
    {
        'title': 'ロボット技術の進化',
        'abstract': 'ロボット技術の進化とその社会への影響について概説する。',
        'summary': 'ロボット技術は、生産性の向上と新たなサービスの提供を可能にする。',
        'url': 'http://example.com/paper8',
        'created_at': '2023-06-08 00:00:00'
    },
    {
        'title': 'IoTの最新動向',
        'abstract': 'インターネット・オブ・シングス（IoT）の最新の動向とその応用について解説する。',
        'summary': 'IoTは、物理的な世界とデジタルな世界をつなげ、新たな価値を創出する。',
        'url': 'http://example.com/paper9',
        'created_at': '2023-06-09 00:00:00'
    },
    {
        'title': 'ビッグデータの活用',
        'abstract': 'ビッグデータの活用方法とそのビジネスへの影響について考察する。',
        'summary': 'ビッグデータは、ビジネスの意思決定を支え、競争優位性を確保するための重要な資源である。',
        'url': 'http://example.com/paper10',
        'created_at': '2023-06-10 00:00:00'
    },
]

for paper in papers:
    paper_id = insert_paper(paper['title'], paper['abstract'], paper['summary'], paper['url'])

試しに実行してみると以下のようなデータを得られました。

query_embedding = generate_embedding('データサイエンス')
results = search_similar_embeddings(query_embedding)
print(results[0])

titles = [row[1] for row in results]
print(titles)

量子コンピュータやブロックチェーンの話題は除外できていることが分かりますね。それっぽく動作しているようです。

試してみた所感

私自身が仮想テーブルの仕組みを分かっておらず、埋め込みベクトル以外のカラムを設定しようとしてエラーでハマったり、そもそも当初はトリガーで自動的に設定されるようにしていたものの、何か上手く動かなくてハマったりと、実際に触ってみると意外とハマりポイントが多くて時間が溶けてしまったのですが、いったんの完成形が分かっていればハマらずに利用できるかなと思います。

埋め込みベクトルもローカルで生成するようにすればもっとコストは抑えられると思いますが、これは未検証です。

SQLiteをフル活用できると、Fly.ioでの運用が捗りますね！

現場からは以上です。