Ruby on Rails & Postgresqlでベクトル近傍検索を手軽に実装するためのTips

2023年8月9日 11:00

こんにちは。YOSHINA SaaS事業部エンジニアの高瀬です。

OpenAIとAzureOpenAIServiceの使い分けについての前回の記事に引き続き、ChatGPTと連携するアプリケーションを開発するために役立ちそうな様々なTipsをご紹介していきます。

本記事では Pgvector と neighbor を連携させて、Ruby on Railsでベクトルの近傍検索を可能にする方法についてご紹介します。Pgvector は Postgresql でベクトル近傍検索を可能にする拡張機能です。neighbor は ActiveRecord から Pgvector を利用してベクトル近傍検索を容易に行えるようにするGemです。これらの機能をシンプルなRuby on Rails + Postgresqlの環境に追加し、ActiveRecordからの近傍検索を試して行こうと思います。

前提となる構成

docker composeプロジェクト内でRuby on Railsアプリケーションを開発する状況を想定しています。RailsがActeveRecord連携するバックエンドDBはPostgresqlを利用します。

以下のようなcomposeファイルで、dbとappサービスがある想定です。

services:
  db:
    image: postgres:15.2
    volumes:
      - db_volume:/var/lib/postgresql/data
    env_file:
      ./.env
    expose:
      - "5432"

  app:
    build:
      context: ./app
    ports:
      - 3000:3000
    working_dir: /app
    volumes:
      - ./app:/app
    command: ['bin/rails', 'server', '-p', '3000', '-b', '0.0.0.0']
    env_file:
      ./.env
    depends_on:
      - db

尚、Ruby on Railsアプリケーションをdocker composeで構築するための手順やノウハウ等は本記事では取り扱いません。Dockerのドキュメントクィックスタート: Compose と Rails や様々な他記事をご参考ください。

以降は既にRailsの初期化とDBの初期化と接続が出来ている状況を想定しています。

pgvectorプラグインの導入

pgvectorとは？

まず初めに、pgvector とは何でしょうか？

pgvectorは、

postgresqlのエクステンション
postgresqlに vector というデータ型を扱えるようにし、近傍探索を可能にする

というものです。

Posgresqlでpgvector拡張を有効化する

今、とにかくさっさと使える状態のdocker imageを探している方はこちら (ankane/pgvector) をpullしてご利用ください。 pgvector公式ページの下の方に Docker における使い方解説もありますので、本章は読み飛ばしても問題ありません。

postgres公式イメージにpgvectorを適用する

公式のpostgres imageには拡張機能としてのpgvectorが含まれていません。pgvectorの Installation に従って、公式のpostgres imageにextension controle fileをinstallします。 docker composeで動かすdbコンテナは公式イメージに対して処理を追加したものになるので、dbサービス用のDockerfileを作ります。

db/Dockerfile

FROM postgres:15.2

# buildに必要な依存関係を入れる
RUN apt-get update && \
    apt-get install -y git make gcc postgresql-server-dev-15

# pgvectorをbuildしてinstall
RUN git clone --branch v0.4.4 https://github.com/pgvector/pgvector.git && \
    cd pgvector && \
    make && \
    make install && \
    cd ../ && rm -rf pgvector

compose.yml(変更箇所)

  db:
    build:
      # composeファイルがある場所にdb/Dockerfileを置いている
      context: ./db

docker compose build でimageをbuild後、pgvectorを有効化します。

dbコンテナにログインしてpsqlコマンドを発行する方法もありますが、RailsコンソールからSQLを発行するのがお手軽です。

con = ActiveRecord::Base.connection
con.execute("CREATE EXTENSION vector;")
   (45.7ms)  CREATE EXTENSION vector;
=> #<PG::Result:0x00007f140d0f9da8 status=PGRES_COMMAND_OK ntuples=0 nfields=0 cmd_tuples=0>

有効化に成功しました。

Postgresqlのバージョンが最新化された直後や、様々な依存関係のバージョンが微妙にかみ合わなくて調整しないといけない時など、自力でextensionを有効化できるノウハウはここ一番で助けになってくれると思います。

neighbor によるActiveRecordからの近傍探索

それでは、ベクトルの近傍検索をActiveRecordから使っていきます。 pgvector-ruby と neighbor の2つのGemが有名ですが、本日は後者を触ってみたいと思います。

neighborのREADME.md には、gemの有効化から基本的な使い方まで分かりやすく説明されております。

早速Gemfileにneighborを追加して

gem 'neighbor'

appコンテナにbundle installして、気持ちを高めていきたいと思います。

docker compose run --rm app bundle install

Vector型のフィールドを含むmodelの作成まで

今回はモデル Item に対し、近傍検索可能なVector型のフィールド embedding を定義します。

migrationファイルをgenerateし

docker compose run --rm app bin/rails g model Item
...
...      
      invoke  active_record
      create    db/migrate/20230730024006_create_items.rb
      create    app/models/item.rb
      invoke    test_unit
      create      test/models/item_test.rb
      create      test/fixtures/items.yml

Itemモデルに embedding フィールドを vector 型として定義しましょう。

OpenAIで現在embeddingに利用できるモデル text-embedding-ada-002 の次元数は 1536 ですので、limit: 1536 オプションを渡します。AzureOpenAIで利用できるモデルであっても、次元数は同じです。

class CreateItems < ActiveRecord::Migration[7.0]
  def change
    create_table :items do |t|
      t.string :name
      t.vector :embedding, limit: 1536 # dimensions

      t.timestamps
    end
  end
end

最後にmigrationを走らせましょう。

docker compose run --rm app bin/rails db:migrate

app/db/schema.rbを確認すると、以下のようなテーブルが CREATE されました。

  create_table "items", force: :cascade do |t|
    t.string "name"
    t.vector "embedding", limit: 1536
    t.datetime "created_at", null: false
    t.datetime "updated_at", null: false
  end

続いて、app/models/item.rb でモデル側には has_neighbors 属性を定義します。

class Item < ApplicationRecord
  has_neighbors :embedding
end

"has_neighbors :embedding" により、近傍検索をするためのメソッド nearest_neighbors が embedding列に対して利用可能になります。

ここでちょっと脇道に逸れて補足的な話ではありますが、この has_neighbors の実装をざっと見ておきたいと思います。

まず、neighborがActiveSupport拡張に追加される時に、Neighbor::Model モジュールが extendされます。このモジュールに定義された has_neighborsメソッドのコードを見ていくと、attribute_name 引数にvector型フィールドのカラム名に相当するシンボル(:enbedding)を受け取っています。他にも、dimension や normalize といったパラメータが用意されています。

has_naighbors メソッド内 class_eval に渡されるブロックでは :nearest_neighbors という scopeメソッドが定義されており、ActiveRecord で近傍検索するための仕組みが作られているようです。

ActiveRecordで近傍検索

では、neighborの威力を体感してみましょう。

ユークリッド距離による近傍検索

以下のように、0->10で、同じ値1536次元からなるvectorを持った Item モデルのレコードを11個作成してみます。

(0..10).each do |val|
  vec = [val.to_f] * 1536
  Item.create!(name: "val_#{val}", embedding: vec)
end

[0.0, 0.0 ...] のvectorを持つItemは name="val_0" としました。ユークリッド距離で計測すると val_0 に近いItemは、val_1, val_2, val_3になることを確認してみます。

vec0 = Item.first
neighbors = Item.nearest_neighbors(:embedding, vec0, distance: "euclidean").first(3).map do |n|
  [n.name, n.neighbor_distance]
end

neighbors
=> [["val_1", 39.191835884530846], ["val_2", 78.38367176906169], ["val_3", 117.57550765359255]]

コサイン距離による近傍検索

次に、コサイン距離による近傍検索をやってみたいと思います。neighborではコサイン類似度ではなく、 コサイン距離 を計算しています。具体的な計算の実装箇所はここにありますね。

ベクトルの角度をイメージしやすいように次元数を2に落としてやってみます。コサイン類似度(コサイン距離)を利用する場合は、vectorに対してL2ノルムによる正規化処理が必須のため、Item モデルの has_naighbors 属性に normalize: true オプションを追加します。

class Item < ApplicationRecord
  has_neighbors :embedding, normalize: true
end

vec_a = [1.0, 2.0], vec_b = [-2.0, 1.0], vec_c = [1, 1] の3つを作ってみます。


             y
             |    [a]    2.0
             |           
 [b]         |    [c]    1.0
             |
―――――――――――――|―――――――――――― x
             |

何故かコードブロックで座標を表現しておりますが、vec_aとvec_bは90度の関係、vec_aとvec_cは90度未満になっている状態です。

vec_a = Item.create!(name: "vec_a", embedding: [1.0, 2.0])
=>
#<Item:0x00007fc2037231d0
 id: 1,
 name: "vec_a",
 embedding: [0.447213595499958, 0.894427190999916]

vec_b = Item.create!(name: "vec_b", embedding: [-2.0, 1.0])
=>
#<Item:0x00007fc205f93980
 id: 2,
 name: "vec_b",
 embedding: [-0.894427190999916, 0.447213595499958]

vec_c = Item.create!(name: "vec_c", embedding: [1.0, 1.0])
=>
#<Item:0x00007fc2034d7958
 id: 3,
 name: "vec_c",
 embedding: [0.7071067811865476, 0.7071067811865476]

dbに保存される時に、vectorは既に正規化されているようですね。

vec_a で近傍検索してみましょう。c, bの順に近いはずです。

neighbors = vec_a.nearest_neighbors(:embedding, distance: "cosine").map do |n|
  [n.name, n.neighbor_distance]
end

neighbors
=> [["vec_c", 0.051316651780405786], ["vec_b", 1.0]]

vec_b で近傍検索してみましょう。a, cの順に近いはずです。

neighbors = vec_b.nearest_neighbors(:embedding, distance: "cosine").map do |n|
  [n.name, n.neighbor_distance]
end

neighbors
=> [["vec_a", 1.0], ["vec_c", 1.3162277827398647]]

正しそうですね。簡単な例でしたが、 ActiveRecord を直接利用した近傍検索が体感出来ました。

終わりに

ChatGPTをはじめとする生成系AIと連携したシステムを検討する際には、AIに対して外部知識を提供するために、Databaseの検索処理を伴う場合があると思います。文書ベクトルを利用した類似文検索は実用性の高い手法であり、ChatGPTを扱うための LangChain 系のライブラリでも使いやすく実装されているものかと思います。

RubyのGemである langchainrb を利用する場合でも、Ruby on Rails & Pgvector の連携周りに知見があると、自分たちのシステムに対してカスタマイズをしやすくなると思います。そして、 langchainrb に関しては、次回の記事で触れていきたいと思います。

レトリバでは、一緒に働く仲間を募集しています。詳細はこちら下のリンクから見られますので、ぜひご覧ください。

https://herp.careers/v1/retrieva