洋食の日記

「だ・である」調ではなく「です・ます」調で書きはじめれば良かったなと後悔してる人のブログです

近似最近傍探索ライブラリHnswlibのRuby bindingを作った

はじめに

Hnswlibは、C++で書かれたHierarchical Navigable Small World graphsによる近似最近傍探索ライブラリである。近似最近棒探索のベンチマークでも上位に登場する。Ruby bindingがなかったので作成した。

hnswlib | RubyGems.org | your community gem host

使い方

インストールは、普通にgemコマンドでインストールできる。外部ライブラリもPythonも必要ない。

$ gem install hnswlib

APIは単順にバインドしたものと、それらをラップしたAnnoyライクなHnswIndexを用意した。 検索インデックスの作成は、以下のようになる。データを追加すれば、それでグラフ構造が内部で作られるので、build_indexみたいなメソッドはない。 データベクトルはRuby Arrayで表す。

require 'hnswlib'

n_features = 40 # 検索インデックスに追加するベクトルの次元数
max_item = 1000 # 検索インデックスに追加するデータ数

t = Hnswlib::HnswIndex.new(n_features: n_features, max_item: max_item)

# ランダムな値からなるベクトルを1000件追加する.
max_item.times do |i|
  v = Array.new(n_features) { rand }
  t.add_item(i, v)
end

# インデックスを保存する.
t.save('foo.ann')

検索は、検索インデックス中のベクトルを番号で指定し、それをクエリとして検索を行うget_nns_by_itemと、クエリとして任意のベクトルを与えて検索を行うget_nns_by_vectorを用意した。

require 'hnswlib'

# 作成した検索インデックスを読み込む.
n_features = 40
max_item = 1000
u = Hnswlib::HnswIndex.new(n_features: n_features, max_item: max_item)
u.load('foo.ann')

# アイテム番号0番に近い10件を検索する.
# 近傍とされるアイテムの番号がRuby Arrayで返る.
neighbor_ids = t.get_nns_by_item(0, 10)

# 任意のベクトルに近い10件を検索する.
# include_distancesにtrueを与えると, 距離も返す.
q = Array.new(n_features) { rand }
neighbor_ids, neighbor_dists = t.get_nns_by_vector(1, 10, include_distances: true)

おわりに

Hnswlibは、ベンチマークではAnnoyよりも良い検索性能を得ているので、Annoyでイマイチなときに使ってみると良いかもしれない。

github.com

もう3ヶ月ぐらいヒドい腱鞘炎に悩まされていて、長時間の作業は難しく、オリジナルなものを作るのはシンドイので、bindingライブラリを作った。