🌊

pg_onnx: PostgreSQLのSQLでテキストEmbeddingと類似検索を実現する

に公開
PostgreSQL
ONNX
LLM
embedding
pgvector
tech

pg_onnxの紹介

概要

pg_onnxは、PostgreSQLでONNX（Open Neural Network Exchange）モデルを直接実行できる拡張機能です。この拡張機能を使用することで、SQLクエリ内でテキストからembedding（ベクトル表現）を生成できるようになります。

pg_onnxの主な特徴：

  • データベース内での機械学習推論: PostgreSQL内で直接ONNXモデルを実行し、外部APIを呼び出す必要がありません
  • 効率的なリソース管理: PostgreSQLはクライアント接続ごとにプロセスを作成しますが、pg_onnxはバックグラウンドワーカーとしてonnxruntime-serverを実行し、ONNXファイルごとに1つのセッションを作成・再利用することで、システムやGPUのメモリ不足を防ぎます
  • pgvectorとの統合: pgvector拡張機能と組み合わせることで、セマンティック検索（意味的類似検索）を実現できます
  • GPUサポート: onnxruntimeのCUDA対応により、GPUを使用した高速推論が可能です (CPUでも推論可能)

ユースケース

  • テキストの類似検索（セマンティックサーチ）
  • レコメンデーションシステム
  • テキスト分類

セットアップ方法

1. 必要なライブラリのインストール

Ubuntu/Debianの場合

# 基本的な依存関係のインストール
sudo apt install cmake libboost-all-dev

ONNX Runtimeのインストール ($HOME/onnxにインストール)

git clone --recursive https://github.com/microsoft/onnxruntime.git

nohup ./build.sh \
  --config Release \
  --build_shared_lib \
  --parallel \
  --cmake_extra_defines CMAKE_INSTALL_PREFIX=$HOME/onnx \
  &> build_onnx.log &
  
# ビルドが完了したら
cd build/Linux/Release
make install

ONNX Runtime Extensionsのインストール ($HOME/onnxにインストール)

embeddingモデルを使用する場合、onnxruntime-extensionsが必要です。特にtokenizerを使用する場合は、SPM（SentencePiece）トークナイザーを有効にしてビルドする必要があります。

git clone --recursive https://github.com/microsoft/onnxruntime-extensions.git
cd onnxruntime-extensions

mkdir build
cd build

cmake .. \
  -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_INSTALL_PREFIX=$HOME/onnx \
  -DONNXRUNTIME_DIR=$HOME/onnx \
  -DOCOS_ENABLE_SPM_TOKENIZER=ON

nohup make &

# ビルドが完了したら
make install

2. pg_onnxのビルドとインストール

  • PG_SRCはPostgreSQLのソースコードのディレクトリを設定
  • DCMAKE_INSTALL_PREFIXはpg_onnxのインストール先を設定
# リポジトリのクローン（サブモジュールを含む）
git clone --recursive https://github.com/kibae/pg_onnx.git
cd pg_onnx

# ビルド
mkdir build && cd build

export PG_SRC=$HOME/src/postgresql-17.6
export ONNXRUNTIME_DIR=$HOME/onnx
export PATH=$HOME/pgsql/bin:$PATH
export LD_LIBRARY_PATH=$HOME/onnx/lib64:$HOME/boost/lib:$LD_LIBRARY_PATH

cmake .. \
  -DCMAKE_BUILD_TYPE=Release \
  -DPostgreSQL_INCLUDE_DIR=$PG_SRC/include \
  -DPostgreSQL_SOURCE_DIR=$PG_SRC \
  -DONNXRUNTIME_DIR=$ONNXRUNTIME_DIR \
  -DONNX_RUNTIME_INCLUDE_DIRS=$ONNXRUNTIME_DIR/include/onnxruntime \
  -DONNX_RUNTIME_LIBRARY_DIRS=$ONNXRUNTIME_DIR/lib64 \
  -DONNX_RUNTIME_LIBRARY=$ONNXRUNTIME_DIR/lib64/libonnxruntime.so \
  -DCMAKE_INSTALL_PREFIX=$HOME/pgsql/pg_onnx

make

# インストール
make install

3. PostgreSQLでの拡張機能の有効化

-- pgvector拡張機能を有効化
CREATE EXTENSION IF NOT EXISTS vector;

-- pg_onnx拡張機能を有効化
CREATE EXTENSION IF NOT EXISTS pg_onnx;

4. ONNXモデルのダウンロードと登録

今回は、多言語対応のembeddingモデル「multilingual-e5-small」のONNX版を使用します。

# Pythonでモデルをダウンロード（Hugging Face Hub経由）
from huggingface_hub import hf_hub_download

repo_id = "oga5/multilingual-e5-small-pg-onnx"

# 必要なモデルファイルをダウンロード
tok_path = hf_hub_download(repo_id=repo_id, filename="tokenizer.onnx")
tte_path = hf_hub_download(repo_id=repo_id, filename="text_to_embedding.onnx")

print(f"Tokenizer: {tok_path}")
print(f"Text-to-Embedding: {tte_path}")

SQLでモデルを登録：

-- トークナイザーモデルの登録
SELECT pg_onnx_import_model(
    'e5-tok',                                          -- モデル名
    'v1',                                               -- バージョン
    pg_read_binary_file('/PATH/tokenizer.onnx')::bytea, -- モデルファイル
    '{"ortextensions_path": "libortextensions.so"}'::jsonb, -- onnxruntime-extensionsのパス
    'e5 tokenizer'                                      -- 説明
);

-- Embeddingモデルの登録
SELECT pg_onnx_import_model(
    'e5-embedding',
    'v1',
    pg_read_binary_file('/PATH/text_to_embedding.onnx')::bytea,
    '{"ortextensions_path": "libortextensions.so"}'::jsonb,
    'e5 text to embedding'
);

5. Embedding生成関数の作成

-- 基本的なEmbedding生成関数
CREATE OR REPLACE FUNCTION e5_embedding(input_text text)
RETURNS vector(384)
AS $$
    SELECT array(
        SELECT jsonb_array_elements_text(
            pg_onnx_execute_session(
                'e5-embedding',
                'v1',
                jsonb_build_object('text', jsonb_build_array(input_text))
            )->'embedding'->0
        )::float
    )::vector(384);
$$
LANGUAGE sql
IMMUTABLE;

-- Passage（文書）用のEmbedding生成関数
-- E5モデルでは"passage: "プレフィックスを付ける
CREATE OR REPLACE FUNCTION e5_embedding_passage(input_text text)
RETURNS vector(384)
AS $$
    SELECT e5_embedding('passage: ' || input_text);
$$
LANGUAGE sql
IMMUTABLE;

-- Query（検索クエリ）用のEmbedding生成関数
-- E5モデルでは"query: "プレフィックスを付ける
CREATE OR REPLACE FUNCTION e5_embedding_query(input_text text)
RETURNS vector(384)
AS $$
    SELECT e5_embedding('query: ' || input_text);
$$
LANGUAGE sql
IMMUTABLE;

6. テーブル作成とテストデータの登録

-- テーブル作成
CREATE TABLE llm_test (
    i integer NOT NULL PRIMARY KEY,
    txt text,
    v vector(384)
);

-- HNSW（Hierarchical Navigable Small World）インデックスの作成
-- vector_ip_opsは内積（Inner Product）を使用
CREATE INDEX llm_test_v_idx ON llm_test USING hnsw (v vector_ip_ops);

-- テストデータの挿入
INSERT INTO llm_test (i, txt) VALUES 
('1', 'Machine learning is a subfield of artificial intelligence'),
('2', 'A database is a system for managing data'),
('3', 'PostgreSQL is a powerful open-source database'),
('4', 'Vector search retrieves results by computing similarity'),
('5', 'ONNX is a standard format for machine learning models'),
('6', 'Natural language processing is a technology for handling text'),
('7', 'Embeddings convert text into vectors'),
('8', 'Cosine similarity measures similarity between vectors'),
('9', 'A tokenizer splits text into tokens'),
('10', 'Transformers are a modern neural network architecture'),
('11', 'SQL is a language for manipulating databases'),
('12', 'Indexes improve query performance'),
('13', 'pgvector is a vector extension for PostgreSQL'),
('14', 'Semantic search retrieves based on meaning'),
('15', 'Neural networks mimic the structure of the brain'),
('16', 'Deep learning uses multi-layer neural networks'),
('17', 'Batch processing handles multiple data at once'),
('18', 'Model inference performs prediction with a trained model'),
('19', 'Fine-tuning adapts an existing model to a specific task'),
('20', 'A cross-encoder evaluates the relevance between two texts');

-- Embeddingの生成と登録
UPDATE llm_test SET v = e5_embedding_passage(txt);

7. 類似検索の実行

multilingual-e5-smallモデルはL2正規化済みのembeddingを出力するため、類似度の計算には**内積（Inner Product）**を使用します。pgvectorでは<#>演算子が負の内積を表すため、距離が小さいほど類似度が高くなります。

-- 類似検索の実行
WITH q AS (
    SELECT 'What is machine learning?' AS query
),
qv AS MATERIALIZED (
    SELECT e5_embedding_query(q.query) AS v FROM q
)
SELECT 
    i, 
    txt, 
    t.v <#> qv.v AS distance
FROM llm_test t, qv
ORDER BY distance
LIMIT 5;

実行結果例：

 i  |                         txt                          |     distance      
----+------------------------------------------------------+-------------------
  1 | Machine learning is a subfield of artificial intelligence | -0.8234567
  5 | ONNX is a standard format for machine learning models     | -0.7891234
 18 | Model inference performs prediction with a trained model  | -0.7654321
 15 | Neural networks mimic the structure of the brain          | -0.7543210
 16 | Deep learning uses multi-layer neural networks            | -0.7432109

応用例：トリガーを使った自動Embedding生成

データ挿入時に自動的にembeddingを生成するトリガーを設定することもできます：

-- トリガー関数の作成
CREATE OR REPLACE FUNCTION auto_generate_embedding()
RETURNS TRIGGER AS $$
BEGIN
    NEW.v := e5_embedding_passage(NEW.txt);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- トリガーの設定
CREATE TRIGGER llm_test_embedding_trigger
    BEFORE INSERT OR UPDATE OF txt
    ON llm_test
    FOR EACH ROW
    EXECUTE FUNCTION auto_generate_embedding();

まとめ

pg_onnxを使用することで、PostgreSQL内で直接機械学習モデルを実行し、テキストのembeddingを生成できます。pgvectorと組み合わせることで、ベクトルによる類似検索システムを構築できます。

主なメリット：

  • SQLでembedding modelを利用したベクトルデータの生成や類似検索が可能
  • 外部APIやサービスへの依存を減らせる
  • データベース内で完結するため、レイテンシーが低い
  • GPUサポートにより高速な推論が可能
  • トリガーなどを使った自動化が容易

参考

GitHubで編集を提案

Discussion