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

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

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環境はサポートされていないようです。

## 公式のデモ

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

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

    ```shell
    pip install -Uq voyager
    ```

1. テストデータの取得

    ```shell
    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
    ```

1. Numbyのロード

    ここからPythonのコードです。動画ではIPythonで実施していました。

    ```python
    import numpy as np
    ```

1. テストデータのロード

    テストデータは[Word2Vec](https://ja.wikipedia.org/wiki/Word2vec)でベクトル化したデータだと思われます。

    ```python
    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:]]
    ```

1. データの確認
    
    ```python
    labels[1]
    ```
    
    > 'the'
    
    ```python
    labels[2]
    ```
    
    > 'of'
    
    ```python
    labels[3]
    ```
    
    > 'and'
    
    ```python
    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)
    

1. Voyagerのインデックスを作成

    ```python
    import voyager
    
    index = voyager.Index(voyager.Space.Cosine, 200)
    ```

1. データを登録
    
    ```python
    len(vectors)
    ```
    
    > 10000
    
    ```python
    index.add_items(vectors)
    ```

1. 登録したデータの確認

    ```python
    labels.index("dog")
    ```
    
    > 1902

1. 検索実行

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

1. 検索結果の確認
    
    ```python
    ids
    ```
    
    > array([1902, 4138, 2602, 2974, 5054], dtype=uint64)
    
    ```python
    distances
    ```
    
    > array([5.3644180e-07, 2.7525765e-01, 3.5191095e-01, 3.7909913e-01,
           3.8402766e-01], dtype=float32)


    ```python
    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.

1. インデックスを保存
    
    ```python
    index.save("demo-index.voy")
    ```

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


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

## Titan Embeddings G1で実施

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

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

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

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

1. Titan Embeddings G1を呼び出す関数の作成
    
    テキストを渡してベクトル化した結果を取得する関数を作成しました。
    
    ```python
    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
    ```

1. テストデータの作成

    ```python
    labels = ["dog", "dogs", "cat", "bird", "breed"]
    
    vectors = list(map(lambda x: embedding(x), labels))
    ```

1. Voyagerインデックスの作成とデータの登録
    
    ```python
    import voyager
    index = voyager.Index(voyager.Space.Cosine, 1536)
    index.add_items(vectors)
    ```

1. 検索実行

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

1. 検索結果の確認
    
    ```python
    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.