Spotify製のvoyagerをTitan Embeddings G1でやってみた

2024年1月13日 · 約7分

moritalous

Maintainer of this blog

お知らせ

過去にQiitaに投稿した内容のアーカイブです。

少し前にこのようなニュースが出てました。

Spotify、最近傍検索ライブラリVoyagerをオープンソース化

https://www.infoq.com/jp/news/2023/12/spotify-ann-voyager/

気になっていたのでAmazon Bedrockのテキスト埋め込みモデルTitan Embeddings G1で試してみました。

情報

公式サイト

https://spotify.github.io/voyager/

GitHub

https://github.com/spotify/voyager

PythonとJavaに対応していて、それぞれでインデックスデータは共有できるようです。また、現時点でWindows環境はサポートされていないようです。

公式のデモ

公式サイトにデモ動画に沿って実施しました。（動画しかないので文字起こししておきます）

ライブラリーのインストール
```
pip install -Uq voyager
```

テストデータの取得

wget https://storage.googleapis.com/embedding-projector/data/word2vec_10000_200d_tensors.bytes
wget https://storage.googleapis.com/embedding-projector/data/word2vec_10000_200d_labels.tsv

Numbyのロード
ここからPythonのコードです。動画ではIPythonで実施していました。
```
import numpy as np
```

テストデータのロード

テストデータはWord2Vecでベクトル化したデータだと思われます。

num_dimentions = 200

with open("word2vec_10000_200d_tensors.bytes", "rb") as f:
  vectors = np.fromfile(f, np.float32).reshape(-1, num_dimentions)

with open("word2vec_10000_200d_labels.tsv", "r") as f:
  labels = [line.split("\t")[0] for line in f.readlines()[1:]]

データの確認
```
labels[1]
```
'the'
```
labels[2]
```
'of'
```
labels[3]
```
'and'
```
vectors[0]
```
array([-0.141771 , 0.249576 , -0.188584 , -0.0815223 , 0.128442 , 0.547185 , -0.197366 , 0.142269 , -0.438946 , -0.0416157 , 0.370258 , 0.408382 , 0.011527 , 0.274481 , -0.00205041, 0.165503 , -0.0883049 , 0.286902 , -0.0418692 , 0.0736817 , -0.0211798 , -0.0613568 , 0.17691 , -0.141145 , 0.0105192 , -0.226281 , -0.324913 , 0.266758 , -0.104392 , -0.170748 , 0.00121182, -0.0421411 , -0.126701 , -0.335706 , -0.0201676 , -0.314706 , 0.227294 , 0.181603 , 0.103264 , 0.333935 , 0.0354471 , 0.0635742 , 0.205139 , 0.249157 , -0.136408 , -0.0435595 , 0.095526 , -0.0772112 , -0.00595369, -0.182302 , 0.323586 , -0.204001 , -0.0916038 , 0.0807845 , -0.243777 , 0.119747 , 0.00691663, -0.1902 , 0.263702 , 0.244449 , 0.0441175 , 0.0958303 , -0.0618864 , -0.202204 , -0.342813 , 0.318309 , -0.094516 , 0.307758 , 0.109799 , 0.179937 , 0.209827 , 0.270957 , 0.0364361 , -0.26446 , 0.364246 , -0.366684 , -0.0671071 , 0.10821 , -0.259401 , 0.0538033 , 0.159056 , -0.206028 , -0.0396725 , -0.336107 , 0.234527 , -0.0116628 , -0.0904525 , -0.212477 , -0.408584 , -0.0243164 , -0.274519 , 0.403208 , 0.215137 , -0.132371 , -0.0714231 , 0.188115 , 0.0488086 , -0.0825052 , -0.0325133 , -0.143736 , -0.0349842 , 0.646822 , 0.17857 , 0.386225 , 0.266737 , 0.261134 , 0.0294623 , 0.069207 , 0.0511157 , 0.145836 , -0.0794004 , 0.204002 , -0.193424 , 0.157486 , -0.0425395 , -0.297959 , -0.0443972 , -0.24584 , 0.328743 , -0.0362118 , -0.109993 , 0.368324 , -0.0865976 , -0.0313383 , 0.148474 , ... 0.183288 , 0.0209599 , -0.104203 , -0.169894 , -0.107271 , 0.292521 , -0.177604 , -0.108201 , -0.367897 , -0.281144 , 0.0879999 , -0.291526 , -0.231764 , 0.17579 , 0.0101314 , 0.161831 , -0.0566941 , -0.0891432 , 0.263995 , -0.303327 ], dtype=float32)

Voyagerのインデックスを作成

import voyager

index = voyager.Index(voyager.Space.Cosine, 200)

データを登録
```
len(vectors)
```
10000
```
index.add_items(vectors)
```
登録したデータの確認
```
labels.index("dog")
```
1902

検索実行

vector = vectors[labels.index("dog")]

ids, distances = index.query(vector, 5)

検索結果の確認
```
ids
```
array([1902, 4138, 2602, 2974, 5054], dtype=uint64)
```
distances
```
array([5.3644180e-07, 2.7525765e-01, 3.5191095e-01, 3.7909913e-01, 3.8402766e-01], dtype=float32)
```
for id, distance in zip(ids, distances):
  print(f"\t{labels[id]!r} is {distance:.2f} away from dog.")
```

>   'dog' is 0.00 away from dog.
    'dogs' is 0.28 away from dog.
    'cat' is 0.35 away from dog.
    'bird' is 0.38 away from dog.
    'breed' is 0.38 away from dog.

インデックスを保存
```
index.save("demo-index.voy")
```
注記
保存したインデックスの読み込みはこのように行います。
index = voyager.Index.load("demo-index.voy")

使い方は理解できたでしょうか？あくまでベクトルの検索に特化しているので、ベクトル化した値から元の文書を取得したい場合は別で管理する必要があります。

Titan Embeddings G1で実施

上記デモと同様の内容をAmazon BedrockのTitan Embeddings G1で行いました。

ライブラリーのインストール
```
pip install -Uq voyager boto3
```

ライブラリーのロードとBedrockクライアントの作成

import json
import boto3
client = boto3.client('bedrock-runtime')

Titan Embeddings G1を呼び出す関数の作成

テキストを渡してベクトル化した結果を取得する関数を作成しました。

def embedding(inputText: str):
  body = {
    "inputText": inputText
  }

  response = client.invoke_model(
    modelId="amazon.titan-embed-text-v1",
    contentType="application/json",
    accept="*/*",
    body=json.dumps(body)
  )

  embedding = json.loads(response["body"].read())["embedding"]
  return embedding

テストデータの作成

labels = ["dog", "dogs", "cat", "bird", "breed"]

vectors = list(map(lambda x: embedding(x), labels))

Voyagerインデックスの作成とデータの登録

import voyager
index = voyager.Index(voyager.Space.Cosine, 1536)
index.add_items(vectors)

検索実行

vector = vectors[labels.index("dog")]
ids, distances = index.query(vector, 5)

検索結果の確認

for id, distance in zip(ids, distances):
  print(f"\t{labels[id]!r} is {distance:.2f} away from dog.")

'dog' is 0.00 away from dog.
'dogs' is 0.14 away from dog.
'cat' is 0.25 away from dog.
'breed' is 0.47 away from dog.
'bird' is 0.48 away from dog.

特にハマるところもなく、使えました。

まとめ

簡単に使えるかつ、Bedrockでの利用も可能でした。

また、事前準備されたテストデータとTitan Embeddings G1で、検索した結果（近さだけでなく順番も）が変わることがわかりました。このあたりがベクトル化の性能ということでしょうか？

Word2Vec

'dog' is 0.00 away from dog. 'dogs' is 0.28 away from dog. 'cat' is 0.35 away from dog. 'bird' is 0.38 away from dog. 'breed' is 0.38 away from dog.

Titan Embeddings G1

'dog' is 0.00 away from dog. 'dogs' is 0.14 away from dog. 'cat' is 0.25 away from dog. 'breed' is 0.47 away from dog. 'bird' is 0.48 away from dog.

情報​

公式のデモ​

Titan Embeddings G1で実施​

まとめ​

情報

公式のデモ

Titan Embeddings G1で実施

まとめ