背景
近年、生成AIの進化が凄まじく新たなモデルやサービスが次々と登場し、ニュースを聞かない日はないですね。それと同時にこれらの生成AIを運用中のサービスに取り入れたいという要望も増えてきているかと思います。
そういった際は、RAG（Retrieval-Augmented Generation）などの手法を用いて、ベクトルデータベースの構築やファインチューニングが必要になるかと思います。
そこで今回はRails/PostgreSQL/pgvectorを組み合わせて、ベクトルデータベースを活用するサンプルを実装してみることにしました。
ゴール
Docker上でRailsのActiveRecordと、後述するPostgreSQLの拡張機能であるpgvectorを連携し、簡単なベクトル検索をしてみようと思います。
またベクトルの埋め込みに必要な値を計算するためのAPIをPythonのFlaskで作成します。
サンプル実装
ベクトルDB
ベクトルDB（Vector Database）は、ベクトルデータを格納、管理、検索するためのデータベースです。
ベクトルデータは、数学的なベクトルで表現され、空間上の位置や方向を表します。例えば、テキストデータをベクトル化する場合、各単語や文章をベクトルとして表現し、それらのベクトル間の距離や関連性を計算することができます。
ベクトルDBにはOSSで公開されているものやローカルで動作するもの、クラウド上で動作するものなど、様々なプロダクトが登場しています。
pgvector
pgvectorはPostgreSQLに、ベクトルDBの機能を提供する拡張機能です。
この拡張機能を使うとPostgre内にベクトルデータを格納し、特定のベクトルと最も近い要素を検索することができます。
既にPostgreSQLで運用しているサービスであれば、新たにサービスを追加する必要がないので比較的に導入しやすいのではないでしょうか。
またCloudSQLやAlloyDB for PostgreSQLを使用している環境下においてもインストールすることが可能です。
環境構築
実際に環境をdocker-compose上に作成していきます。
まずはpostgressとRailsを用意します。
services:
postgres:
image: pgvector/pgvector:pg16
environment:
POSTGRES_HOST: postgres
POSTGRES_DB: pgvector_dev
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
volumes:
- postgres_volume:/var/lib/postgresql/data
ports:
- "5432:5432"
web:
build: .
command: bash -c "rm -f tmp/pids/server.pid && bundle exec rails s -p 3000 -b '0.0.0.0'"
volumes:
- .:/myapp
ports:
- "3000:3000"
tty: true
depends_on:
- postgres
volumes:
postgres_volume:
PostgreSQLのコンテナについて、すでにpgvectorが適用されたDockerが公開されているため、そちらを利用しました。
別のPostgreSQLバージョンを使いたい場合は拡張機能としてpgvectorをインストールし有効にする必要があります。
neighborのインストール
neighborはActiveRecordからpgvectorを利用したベクトルデータを操作可能にするためのGemです。
Gemfileに以下を追加し、マイグレーションを実行します。
gem "neighbor"
neighborにはcubeとvectorの２つのモードが存在します。
cubeはPostgreSQL 9.5から使えるようになった多次元立体を表すためのcubeデータ型です。
今回はエンベディングを保存したいので、vectorモードでmigrateします。
# rails generate neighbor:vector
create db/migrate/20240409050025_install_neighbor_vector.rb
実行するとvectorを有効化するためのマイグレーションファイルが追加されています。
class InstallNeighborVector < ActiveRecord::Migration[7.0]
def change
enable_extension "vector"
end
end
モデルの作成
ベクトル検索の対象となるモデルを作成します。
今回はWeb上に公開された記事を格納することを想定したItemという名前のモデルを作ることにします。
bodyに入っているテキストをエンベディングし検索できるようにする想定です。
# rails generate model Item url:string date:datetime body:text
invoke active_record
create db/migrate/20240411214904_create_items.rb
create app/models/item.rb
invoke test_unit
create test/models/item_test.rb
create test/fixtures/items.yml
# rails db:migrate
== 20240411214904 CreateItems: migrating ======================================
-- create_table(:items)
-> 0.0042s
== 20240411214904 CreateItems: migrated (0.0042s) =============================
さらにItemモデルに対してエンべディングされた値を保存するカラムをvector型で追加するマイグレーションファイルを作成します。
# rails g migration AddEmbeddingToItems
invoke active_record
create db/migrate/20240411215904_add_embedding_to_items.rb
作成されたマイグレーションファイルは以下のような内容になっています。
class AddEmbeddingToItems < ActiveRecord::Migration[7.0]
def change
add_column :items, :embedding, :vector, limit: 1024 # dimensions
end
end
ここで指定しているlimitの値はエンべディングで使用するモデルの次元数を指定します。
OpenAIが使用している「text-embedding-ada-002」であれば1536となります。
今回は日本語が含まれているので多言語用モデルの「Multilingual-E5」を使用することにしました。
こちらは一番精度が高いlargeの場合の次元数は1024なのでそちらに合わせます。
ローカルで動作させる場合は計算が早いsmallを利用します。
次にneighborによって追加されたメソッドをモデルから使えるようにするため、has_neighborsを指定します。
class Item < ApplicationRecord
has_neighbors :embedding
end
以上でRailsとpgvectorの準備が完了しました。
Embedding
ここでのEmbeddingとは単語や文章を、その意味を表現するベクトル空間に配置するために、計算可能な形に変換することを意味します。
Embeddingの方法は様々あり、一般的にはPythonのsentencetransfomerを使う方法が手軽かと思います。
その他にもOpenAIのEmbeddingsAPIで取得したり、Ollamaなどのローカルで動作するLLMから取得する方法も存在します。
OpenAIのEmbeddings等を利用すると簡単に利用できるのですが、当然利用料金が発生してくるため、今回はローカルで動作するEmbedding用APIをFlaskで作ることにしました。
Embedding用APIの準備
Flask用のDockerfileを作成します。
FROM python:3.10
WORKDIR /app
ENV FLASK_APP=app
COPY ./requirements.txt ./
COPY ./app.py ./
RUN pip install --upgrade pip
RUN pip install -r requirements.txt
docker-composeに以下を追加します。
flask:
build:
context: ./flask/
dockerfile: Dockerfile
ports:
- "5000:5000"
container_name: flask
volumes:
- ./flask:/usr/src/app
command: flask run --host=0.0.0.0
jsonで受け取った文章をembeddingして返却するだけのAPIを作成します。
from flask import Flask, jsonify, request
from sentence_transformers import SentenceTransformer
app = Flask(__name__)
@app.route("/")
def index():
return "index page"
@app.route('/api/embeddings', methods=['POST'])
def api_embeddings():
content = request.json
sentences = content["data"]
embeddings = SentenceTransformer('intfloat/multilingual-e5-small').encode(sentences)
return jsonify({"embeddings": embeddings.tolist()})
次にRailsからEmbeddingAPIへリクエストを送るクライアントを作成します。
class FlaskClient
require 'net/http'
FLASK_API_URL = 'http://flask:5000/api/embeddings'
def self.embeddings(data)
res = Net::HTTP.post(
URI(FLASK_API_URL),
{ data: }.to_json,
'Content-Type' => 'application/json'
)
JSON.parse(res.body)['embeddings']
end
end
レコードの更新時に上記クライアントを使ってembeddingの値をセットするようにしておきます。
class Item < ApplicationRecord
has_neighbors :embedding
before_save -> { self.embedding = FlaskClient.embeddings(body) }
end
これでEmbeddingの値を取得、保存する準備ができました。
サンプルデータセットの登録
次に動作確認をするためのサンプルデータを取得します。
今回はLivedoorニュースのデータセットを使います。
こちらは株式会社ロンウイットさんが、収集、公開してくださっているデータです。
ダウンロート、解凍後は以下のスクリプトでモデルを作成していきます。
今回はデータセットの中でもITニュース記事のみを取り込むことにしました。
Dir.glob('./lib/tasks/text/it-life-hack/it-life-hack-*.txt').each do |file_path|
File.open(file_path, 'r') do
content = _1.read.split("\n")
Item.create!(url: content[0], date: content[1], body: content[2..].join)
end
end
これで850件程のサンプルデータを登録することができました。
データの検索
それでは実際にデータが登録できたのでneighborによる検索を検証してみます。
どれでもいいので検証用のレコードを取得します。
今回はアップルに関する記事を取得しました。
irb(main):049> item
=>
#<Item:0x0000ffffaae1d138
id: 7,
url: "http://news.livedoor.com/article/detail/6294340/",
date: Mon, 20 Feb 2012 00:00:00.000000000 UTC +00:00,
body:
"アップル、デベロッパプレビューをリリース！次期Mac OS X「Mountain Lion」が明らかにアップルは2012年2月16日（米国カリフォルニア州クパティーノ現地時間）、9番目のメジャーリリースとなる「OS X Mounta
in Lion」のデベロッパプレビューをリリースした。「iPadの人気アプリケーションや機能をMacにもたらし、OS Xのイノベーションを加速させるもの」としているが、どこが凄いのだろうか。MountainLionはメッセージ、メモ、リマ...",
created_at: Thu, 25 Apr 2024 21:23:45.286110000 UTC +00:00,
updated_at: Thu, 25 Apr 2024 21:23:45.286110000 UTC +00:00,
次にこの記事に類似した記事を5件とってみます。
irb(main):075> item.nearest_neighbors(:embedding, distance: "euclidean").first(5).pluck(:body).map{_1.truncate(100)}
=>
["Macの次期OS「Mountain Lion」の真の狙いとは？ 【役立つセキュリティ】先日、Appleが「Mac OS X Mountain Lion Developer Preview」をリリ...",
"Mountain Lion登場！ 25日よりMac App StoreからMountain Lion公開7月25日、Appleの最新OSとなる「OS X Mountain Lion」が正式公開。...",
"iPhone5やiPad miniは姿を見せるか？アップルのWWDCを日本語で解説アップルは2012年6月11日〜6月15日（現地時間）、サンフランシスコにあるMoscone Center We...",
"新型MBP登場で盛り上がったWWDC！ iPhone 5も7インチiPad miniもなし！「iPhone5やiPad miniは姿を見せるか？アップルのWWDCを日本語で解説」を見た人には、い...",
"新型iPad情報でビンゴも！ITのいまを伝える【デジ通】まとめ読みIT系の情報に鋭く切り込んでいたかと思えば、トホホなレビュー連載記事を上中下の3回も続けてしまったりと、波乱万丈な【デジ通】によ..."]
アップルに関する類似記事が取得できました。
has_neighborsを指定したモデルからはnearest_neighborsというメソッドを使うことができます。
こちらは名前のとおり特定のインスタンスから一番近いレコードを取得することができます。
呼び出し時は埋め込みに使ったモデルのカラム名と、distanceを指定する必要があり、ユークリッド距離（euclidean）とコサイン類似度(cosine)がサポートされています。
また実際に検索の起点にしたいvecotorの位置が決まっている場合は、以下のように実際の値を指定して取得することができます。
Item.nearest_neighbors(:embedding, [0.9, 1.3, 1.1], distance: "euclidean").first(5)
今回はMultilingual-E5のモデルをsmallにしていますが、largeにすることでさらに精度の向上が見込めると思います。
インデックス
何も指定しないと線形探索によるフルスキャンになってしまいます。
pgvectorではHNSWとIVFFlatの2種類のインデックスをサポートしています。
class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.1]
def change
add_index :items, :embedding, using: :hnsw, opclass: :vector_l2_ops
# or
add_index :items, :embedding, using: :ivfflat, opclass: :vector_l2_ops
end
end
IVFFlatは、ベクトルをリストに分離し、クエリベクトルに最も近いリストの選択したサブセットを検索します。
HNSWは、ベクトルを多層のグラフ構造に並び替えて検索します。
どちらも精度を犠牲にしてその分速さを出すことになります。
IVFFlatの方がメモリ使用量は少ないのですが、ある程度データがないとインデックスが効かないようです。
HNSWの方が後発で精度も高く、データがなくてもインデックスを効かせることができるのですが、IVFFlatに比べるとインデックス構築時間が長く、メモリ使用量は多くなるそうです。
終わりに
今回はRails、PostgreSQL、pgvectorを使ってベクトル検索を試しました。
また、その過程でテキストの埋め込みに関しては、FlaskのAPIを作成してみました。
ローカルでの検証時に、毎回OpenAIのEmbeddingsなどを使用するのに比べるとコスト削減になるのではないでしょうか。
この他にもベクトル検索の手法は多数ありますが、既にRailsとPostgreSQLで運用しているシステムにベクトル近傍検索を組み込む場合の選択肢としては、pgvectorを検討してもよいかと思いました。
また、この手法を応用することで、生成AIのRAGとして使うこともできるのでLangChainなどとの連携も簡単にできるのではないかと考えています。
Discussion