Azure OpenAI Serviceと仮想マシン上のRedisで埋め込みとベクトル検索してみる

Azure openAI ServiceのEmbeddingsを用いて今流行りのベクトル検索をやってみた。という内容です。
SaaSのベクトルデータベースを使わずになんとかならないかなーと思ったのがきっかけです。

方法として、Redisを使うのですがRediSearchというモジュールがベクトル検索を行うのに必要で、PaasのAzure Cache for RedisではEnterpriseプランでしか利用できないとのこと。

えええ…！このEnterpriseプランおったまげるほど高いぞ…！

一番安いE10プランでも月13万くらいする…

しかし、Redis Stack Serverを構築すればイケるという事なので、Azure仮想マシンでRedis環境を構築する事にしました。

今回下記の記事をとても参考にさせて頂いています。

1 Azure OpenAI ServiceでEmbeddingsのモデルを作成

text-embedding-ada-002モデルをデプロイします。
Azure OpenAI Serviceは申請が必要なので、本家のOpenAIでもOK。

キーとデプロイしたモデルの名前をメモっておく。

2 データセットに埋め込みを生成する

分割してコードを記載していきます。
まず、これから使うライブラリをインポート

python

import os
import ast
from dotenv import load_dotenv
import numpy as np
import openai
import pandas as pd
import redis
import tiktoken
from openai.embeddings_utils import get_embedding, cosine_similarity
from redis.commands.search.query import Query, Filter
from redis.commands.search.result import Result
from redis.commands.search.field import VectorField, TextField

.envファイルに先程のキーとモデル名を書いておきます。

env

OPENAI_NAME=<モデル名>
OPENAI_KEY=<アクセスキー>

OpenAIの初期設定もろもろ。

python

load_dotenv()

openai_name = os.environ.get("OPENAI_NAME")
openai_uri = f"https://{openai_name}.openai.azure.com/"

openai.api_type = "azure"
openai.api_base = openai_uri
openai.api_version = "2022-12-01"
openai.api_key = os.environ.get("OPENAI_KEY")

# embedding_model_for_[doc/query] = "デプロイしたモデルの名前" 
embedding_model_for_doc = "text-search-doc"
embedding_model_for_query = "text-search-query"
embedding_encoding = "gpt2"
max_tokens = 2000

■検索用で使うデータセット

今回つかう検索用のデータセットは以前使った居酒屋さんのメニュー表的なやつを使います。

csv

項目,タイトル,価格
焼酎,（芋）大隅、（芋）金黒,530円,
焼酎,（芋）農家の嫁,630円,
焼酎,（麦）村正,560円,
焼酎,（麦）吟麗玄海,580円,
焼酎,（米）もっこす,560円,
焼酎,（米）鳥飼,680円,
焼酎,（泡盛）三年古酒 瑞泉,800円,
日本酒,松竹梅 上撰 豪快（兵庫）,720円,
日本酒,越乃景虎 本醸造 超辛口（新潟）,"1,080円",
...

埋め込みを実行していきます。

python

input_datapath = "data/data.csv"
df = pd.read_csv(input_datapath)

#タイトルと価格をあわせてCombined列を作る
df["Combined"] = (
    "title: " + df.title.str.strip() + "; Cost: " + df.cost.str.strip()
)

encoding = tiktoken.get_encoding(embedding_encoding)

# 埋め込み実行 data_embeddings.csvとして保存しておく。
df["Embedding"] = df["Combined"].apply(lambda x: get_embedding(x, engine=embedding_model_for_doc))
df.to_csv("data/data_embeddings.csv")

python

#できてるか確認
df["Embedding"][0]

[0.009056703187525272,
 -0.01744893752038479,
 -0.002270820550620556,
 0.005382193252444267,
 0.004461904522031546,
 -0.007069943007081747,
 -0.009907222352921963,
 -0.026645179837942123,
 -0.0163459200412035,
 -0.022964024916291237,
...

大丈夫そう。
生成したベクトルの次元数を確認しておきます。

python

len(df["Embedding"][0])

# ベクトル空間の次元数は1536
embedding_dimension = 1536  

3 仮想マシンにRedis Stack Serverを構築

下記のドキュメントを参考におこないます。

Azureでの仮想マシン作成手順は割愛します。
OSはubuntu 20.04で検証しました。
Dockerで構築するのが早そうなので、公式の通りに進めます。

まずはDockerをインストール

bash

sudo apt-get update
sudo apt-get install apt-transport-https ca-certificates curl software-properties-common

#Dockerの公式GPGキーを追加
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg

#Dockerリポジトリを追加。
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

#Dockerリポジトリを更新して、Dockerパッケージをインストール。
sudo apt-get update
sudo apt-get install docker-ce docker-ce-cli containerd.io

bash

#確認
sudo docker --version
Docker version 23.0.5, build bc4487a

bash

# redis-cliを使うためにインストールしておく
sudo apt install redis-tools

DockerでRedisStackサーバーを起動します。
起動方法オプションについては公式を参照。

bash

sudo docker run -v ~/local-data/:/data -d --name redis-stack-server -p 6379:6379 redis/redis-stack-server:latest

Redisに接続して確認

bash

redis-cli -h 127.0.0.1 -p 6379
127.0.0.1:6379>

bash

127.0.0.1:6379>  FT._LIST
(empty list or set)

indexは作成してないので何も無い状態。

4 ベクトル検索してみる

ローカルから仮想マシン上のRedisに接続できるように一旦ポート開放しておく。

検証目的なので、試し終わったら仮想マシンごと消します。

python

#hostは仮想マシンのパブリックIP
redis_conn = redis.StrictRedis(host="**.***.***.***",port=6379)

python

# インデックスを作成する。
schema = ([
    TextField("itemid"),
    VectorField("Embedding", "HNSW", {"TYPE": "FLOAT32", "DIM": embedding_dimension, "DISTANCE_METRIC": "COSINE" }),
    TextField("title"),
    TextField("cost"),
    TextField("Combined"),
])

redis_conn.ft().create_index(schema)

仮想マシンのRedisを確認するとインデックスが作成されている事が分かる。

127.0.0.1:6379> FT._LIST
1) "idx"

Redisにデータを追加していく

python

# Vm上のRedisにdfのデータを追加
def load_vectors(client,df):
    p = client.pipeline(transaction=False)

    for i in df.index:
        key = f"doc:{str(i)}"
        data = {
            "itemid": str(i),
            "Embedding": np.array(df["Embedding"][i]).astype(np.float32).tobytes(),
            "title": df["title"][i],
            "cost": df["cost"][i],
            "Combined": df["Combined"][i],
        }
        # データを追加する
        client.hset(key, mapping=data)

    p.execute()

load_vectors(redis_conn, df)

python

#データが追加されたか確認する
print("Index size: ", redis_conn.ft().info()['num_docs'])

Index size:  119

ちゃんと追加されたようです。

一応、Redisのほうでも確認

127.0.0.1:6379> KEYS *
  1) "doc:97"
  2) "doc:115"
  3) "doc:18"
  4) "doc:31"
  5) "doc:103"
  6) "doc:19"

127.0.0.1:6379> FT.INFO idx
 1) index_name
 2) idx
 3) index_options
...
 9) num_docs
10) "119"

num_docsが119なのでインデックスにも登録されている。OKそう。

python

# クエリにマッチするメニューを検索して返す関数
def search_reviews_redis(query, n=3, pprint=True, engine=embedding_model_for_query):
    q_vec = np.array(get_embedding(query, engine=engine)).astype(np.float32).tobytes()
    
    q = Query(f"*=>[KNN {n} @Embedding $vec_param AS vector_score]").sort_by("vector_score").paging(0,n).return_fields("vector_score", "Combined").return_fields("vector_score").dialect(2)
    
    params_dict = {"vec_param": q_vec}
    ret_redis = redis_conn.ft().search(q, query_params = params_dict)
    
    columns = ["Similarity", "Ret_Combined"]
    ret_df = pd.DataFrame(columns=columns)

    for doc in ret_redis.docs:
        # コサイン距離をコサイン類似度に変換
        sim = 1 - float(doc.vector_score)

        com = doc.Combined[:200].replace("title: ", "").replace("; cost:", ": ")

        append_df = pd.DataFrame(data=[[sim, com]], columns=columns)

        ret_df = pd.concat([ret_df, append_df], ignore_index=True, axis=0)

        if pprint:
            print("%s | %s\n" % (sim, com))
    return ret_df

検索実行してみます。

results_redis = search_reviews_redis("くたびれたサラリーマンが好むお酒")

結果

0.01746213436099997 | title: 雁木 純米 無濾過生原酒（山口）; Cost: 1,080円

0.005576133728000032 | title: 純米超辛口 船中八策（高知）; Cost: 1,280円

0.003577470778999947 | title: 酔鯨 吟麗 純米吟醸（高知）; Cost: 1,280円

どうなの？誰かおしえてお酒詳しいひと🍺

ちなみに、クエリの書き方で検索文字を含むとするとこうなる。

python

q = Query(f"@title:(*{query}*)=>[KNN {n} @Embedding $vec_param AS vector_score]").sort_by("vector_score").paging(0,n).return_fields("vector_score", "Combined").return_fields("vector_score").dialect(2)

公式に色々なクエリの書き方が記載あるので色々試してみようと思います。

さいごに

まだまだ課題というか、あれできないのかなー？これできないのかなー？とか「あれれーおかしいなー」みたいな部分はあるんですが、とりあえず仮想マシン上のRedisで動く事がわかりました！

要件によってはSaaSサービスが使えないときあると思うので、そういった時に重宝しそうです。
うーん、確かにAzure Cognitive Searchでベクトル検索ができれば万事解決感はあります😅

近々、SaaSサービスのPineConeも試してみよう…

使い終わったAzureリソースは削除することを忘れずにっ！

参考にさせて頂いた記事