Azure OpenAI Service と Azure Cache for Redis でベクトル検索を行う

はじめに

タイトルのとおり、Azure OpenAI Service と Azure Cache for Redis を使ってベクトル検索を行うための一連の方法について、参考にした情報へのリンクと合わせてまとめました。
前半では、Python スクリプトから Azure OpenAI Service の Embeddings (埋め込み) モデルを呼び出してテキストの埋め込み (ベクトル生成) を行い、ローカル PC 上で検索を行います。
後半では、同じく Python スクリプトから、前半で生成したベクトルを Azure Cache for Redis 上に展開して RediSearch モジュール に含まれる Vector Similarity を使って検索を行います。

なお、Azure OpenAI Service 自体については過去の記事でまとめています。
https://zenn.dev/ryo117/articles/1a15305021cd01

コード

本記事で使用したコードは下記で公開しています。

• rsk800/openai-examples/embedding/semantic-text-search-using-azure-openai-and-redis-e2e.ipynb

作りとしては下記の OpenAI や Redis のサンプルコードをベースにしています。一連の作業を End-to-End (通し) で行うように組みなおして、さらに Azrue OpenAI Service および Azure Cache for Redis Enterprise 向けに書き直しています。

• openai/openai-cookbook/examples/Obtain_dataset.ipynb
• openai/openai-cookbook/examples/Semantic_text_search_using_embeddings.ipynb
• RedisAI/vecsim-demo/SemanticSearch1k.ipynb
• Vector Similarity

方法

0. 前提

• Azure のサブスクリプションが存在している。
• リソースグループが作成済み。
• Python 開発環境が構築済み。
  • 参考: 本記事執筆時は Windows 11 Enterprise WSL2 上の Ubuntu 22.04.1 LTS と Python 3.10.6 の組み合わせで検証実施。
• Python 環境に下記ライブラリがインストール済み。
  • numpy
  • openai
  • pandas
  • redis
  • tiktoken

1. Azure OpenAI Serivce による埋め込み

1.1. Azure OpenAI Service リソース作成

こちらを参考にしてリソースを作成します。今回利用する Text Search Embeddings シリーズのモデルは執筆時点で 米国中南部 (South Central US) リージョンと 西ヨーロッパ (West Europe) リージョンでのみ利用可能です。今回はすべてのリソースを 米国中南部 リージョンに作成していきます。

1.2. モデルデプロイ

Azure OpenAI Service が利用可能になったらモデルのデプロイを行います。本記事では Text Search Embeddings シリーズのモデルの中から最も軽量な text-search-ada-doc-001 と text-search-ada-query-001 をそれぞれデプロイします。デプロイは Azure OpenAI Studio と Azure Portal から行うことができます。

【追記】Embedding に関しては Similarity、Text Search、Code Search の全てで V2 モデル (text-embedding-ada-002) の使用が推奨になり、以前の V1 モデルは非推奨になりました。

参考

• 3. Embeddings シリーズ
• Second-generation models

Azure OpenAI Studio から

Azure Portal から

#### 補足情報
Text Search Embeddings シリーズにはモデルの性能別に 4 種類（Davinci、Curie、Babbage、Ada）が存在しています。これらのモデルはそれぞれ埋め込みを行った際に出力されるベクトルの次元数が異なり、表現力と推論時間のトレードオフが存在します。そのため、現実のプロジェクトでは検索結果の正確性や利用シナリオにおける応答時間の要件に応じて最適なモデルを検討していくことになるかと思います。
また、本記事のように出力したベクトルどうしをコサイン類似度により比較する場合、ベクトル空間の次元数によってコサイン類似度の分布が変わるようです。そのため、例えば Davinci で生成したベクトル間のコサイン類似度と Ada で生成したベクトル間のコサイン類似度は単純に比較できないことに注意します。 (e.g. 例えばコサイン類似度が同じ 0.35 だったとしても同じくらい似ているとは言えない)

参考

• コサイン類似度が高いベクトルはどれくらい似ているか(岩波データサイエンス刊行イベントより)

1.3. 埋め込み

1.3.1. 環境変数の設定

Azure OpenAI Service リソース名と API キーを事前に環境変数に登録しておきます。

• 6.2. API キーの取得

# For Ubuntu
export OPENAI_NAME=<your-openai-name>
export OPENAI_KEY=<your-openai-key>

# For Windows
set OPENAI_NAME=<your-openai-name>
set OPENAI_KEY=<your-openai-key>

1.3.2. データ準備

こちらから今回使用するデータセットをダウンロードします。実行する Python スクリプトもしくは Jupyter Notebook と同じディレクトリに data という名前のフォルダを作成して、ダウンロードした CSV ファイルを配置します。内容としては、とある EC サイトにおける商品レビューデータで、言語は英語です。

参考

• オリジナルのデータセット: Amazon Fine Food Reviews

1.3.3. 埋め込み実行

コードの下記部分まで実行するとデータの取り込み、埋め込み、CSV ファイルへの書き出し (途中から再実行するためのバックアップ) が行われます。モデルのパラメーター設定部分では、事前にデプロイした デプロイ名 を設定します。なお、import部分では本記事でこの先使用するライブラリすべてをインポートしています。

import os

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
from redis.commands.search.result import Result
from redis.commands.search.field import VectorField, TextField, NumericField

# Azure OpenAI Service のパラメーター
openai_name = os.environ["OPENAI_NAME"]
openai_uri = f"https://{openai_name}.openai.azure.com/"

openai.api_type = "azure"
# openai.api_base = "https://<your-openai-name>.openai.azure.com/"
openai.api_base = openai_uri
openai.api_version = "2022-12-01"
# openai.api_key = "<your-openai-key>"
openai.api_key = os.environ["OPENAI_KEY"]

# embeddings モデルのパラメーター
# embedding_model_for_doc = "<your-deployment-name>" 
embedding_model_for_doc = "text-search-ada-doc-001"
# embedding_model_for_query = "<your-deployment-name>" 
embedding_model_for_query = "text-search-ada-query-001"
# embedding_encoding = "cl100k_base"
embedding_encoding = "gpt2" # 今回使用するモデルは GPT-2/GPT-3 トークナイザーを使用する
max_tokens = 2000  # 最大トークン数は 2046 のため少し余裕を持った最大値を設定する
embedding_dimension = 1024  # 出力ベクトル空間の次元数は 1024

# データのロードと確認
input_datapath = "data/fine_food_reviews_1k.csv"

df = pd.read_csv(input_datapath, index_col=0)
df = df[["Time", "ProductId", "UserId", "Score", "Summary", "Text"]]
df = df.dropna()
df["Combined"] = (
    "Title: " + df.Summary.str.strip() + "; Content: " + df.Text.str.strip()
)

# ロードしたデータの中から 1000 件をピックアップ
top_n = 1000
df = df.sort_values("Time").tail(top_n * 2)  # 多少減るのを見越して 2000 件サンプリング
df.drop("Time", axis=1, inplace=True)

encoding = tiktoken.get_encoding(embedding_encoding)

# 長すぎるサンプルを除外
df["N_tokens"] = df["Combined"].apply(lambda x: len(encoding.encode(x)))
df = df[df["N_tokens"] <= max_tokens].tail(top_n)

# 埋め込みには数分かかる
df["Embedding"] = df["Combined"].apply(lambda x: get_embedding(x, engine=embedding_model_for_doc))
df.to_csv("data/fine_food_reviews_with_embeddings_1k.csv")

2. テキスト検索 (On ローカル PC)

前のセクションで生成したベクトルを使って検索を行います。

2.1. 検索用関数の定義

検索を実行する関数を定義します。検索クエリのテキストを与えると、クエリのベクトル化を行い、クエリから生成したベクトルと商品レビューから生成したベクトルの間のコサイン類似度を総当たりで計算して最も類似度が高い商品レビューをデフォルトで 3 件返します。ここでは openai ライブラリに含まれる cosine_similarity を使って計算しています。

# クエリにマッチする商品を検索して返す (ローカル PC 上)
def search_reviews(df, description, n=3, pprint=True, engine="text-search-ada-query-001"):
    embedding = get_embedding(description, engine=engine)
    df["Similarity"] = df["Embedding"].apply(lambda x: cosine_similarity(x, embedding))  # コサイン類似度
    df["Ret_Combined"] = df["Combined"].str.replace("Title: ", "").str.replace("; Content:", ": ")
    results = (
        df.sort_values("Similarity", ascending=False)
        .head(n)
        .loc[:,["Similarity", "Ret_Combined"]]
    )
    if pprint:
        for i,r in results.iterrows():
            print("%s | %s\n" % (r["Similarity"], r["Ret_Combined"][:200]))
    return results

2.2. 検索実行

検索を実行します。下記の例では "delicious beans (おいしいお豆)" で検索しています。

results = search_reviews(df, "delicious beans", n=3, engine=embedding_model_for_query)

成功すると下記のような結果が帰ってきます。ここでは表示していませんが、今回用いたデータには商品 ID が含まれていますので、現実的な用途を想定すると検索クエリに最も近い商品を返すといった使い方もできるかと思います。

0.3770173260434461 | Best beans your money can buy:  These are, hands down, the best jelly beans on the market.  There isn't a gross one in the bunch and each of them has an intense, delicious flavor.  Though I hesitate t

0.37512639635782036 | Delicious!:  I enjoy this white beans seasoning, it gives a rich flavor to the beans I just love it, my mother in law didn't know about this Zatarain's brand and now she is traying different seasoning

0.37370195283798296 | Jamaican Blue beans:  Excellent coffee bean for roasting. Our family just purchased another 5 pounds for more roasting. Plenty of flavor and mild on acidity when roasted to a dark brown bean and befor

補足情報

上記で行っている検索は抽象的な表現をした文のクエリを与えても結果を返します。

results = search_reviews(df, "I want exotic and spicy food", n=3, engine=embedding_model_for_query)

結果

0.3373751011799487 | Beautiful:  I don't plan to grind these, have plenty other peppers for that.  I got these to serve whole, especially for my brother who loves spicy Asian cuisine.  And these actually came all the way 

0.33712935581633535 | spicy:  It is a too spicy grocery in japan.<br /><br />If you cook for udon or something, you can use one.<br /><br />You should buy one.

0.3345656673719058 | Great flavor - spicy but not too hot:  I got some of this from a low sodium foods website. I can't believe how much flavor this has for only 170mg per oz.  I am on a strict low sodium diet and I don't

また、GPT-3 ベースのモデルは多言語対応していますので、検索対象と異なる言語でクエリを投げても (精度が良いかは別として) 結果を返します。

results = search_reviews(df, "おいしいお豆", n=3, engine=embedding_model_for_query)

結果

0.3027765950820416 | Simple and Authentic:  This is a fantastic do-it-yourself poke product. Just add sesame oil and green onion for color then enjoy your authentic Hawaiian treat!

0.30150749900539286 | spicy:  It is a too spicy grocery in japan.<br /><br />If you cook for udon or something, you can use one.<br /><br />You should buy one.

0.2937023705102833 | sesamiOil:  This is a good grocery for us.<br /><br />If you cook something,you can use it.<br /><br />It is smells so good.<br /><br />You should buy it.

3. ベクトルデータの Azure Cache for Redis への展開

ここまでは、Azure OpenAI Service を使ってテキストをベクトル化してローカル PC 上で検索を行いました。対象データの件数が少ないうちはこれでもよいのですが、「検索対象データの件数が多い」 and/or 「Davinci など出力されるベクトルの次元が多いモデルを使う」 場合はローカル PC (もしくは単一 VM) での検索はどこかで限界を迎え、そもそも全件分のベクトルがメモリ上に乗りきらなかったり、計算量が増えて検索がいつまでたっても終わらないといった事態が起こり得ます。
そこで、ここからは将来想定されるこれらのボトルネックを解消すべく、スケーラビリティの高い Azure Cache for Redis を使ったベクトルの蓄積・検索を行っていきます。

3.1. Azure Cache for Redis リソース作成

下記ドキュメントの手順を参考にして、Azure Cache for Redis リソースを作成します。

• クイックスタート: Redis Enterprise キャッシュを作成する

Azure Portal の Marketplace から Redis Cache を検索します。

リージョンは Azure OpenAI Service と同じ米国中南部 (South Central US)リージョンを選択します。ベクトルデータの蓄積と検索を行う RediSearch モジュールは Enterprise Tier のみで利用可能なため、キャッシュの種類は Enterprise のいずれかを選択します。今回は Enterprise の中で最も安い Enterprise E10 を選択します。

詳細タブのモジュールからRediSearchのチェックを入れます。また、クラスタリングポリシーは Enterprise を選択します。

Azure Cache for Redis Enterprise の価格について

Azure Cache for Redis Enterprise は執筆時点の価格で米国中南部リージョンの場合、最も小さい SKU / ノード数を選択しても月額 $1,000 近く (1時間あたり $1.336) かかります。キャッシュの停止はできず、キャッシュ上に展開しているデータサイズに関わらずこの金額がかかり続けるため、検証目的で使用する場合はコスト節約のため検証後すぐにリソースごと削除します。

参考

• Azure Cache for Redis での Redis モジュールの使用
• キャッシュへのモジュールの追加
• サービス階層
• 機能の比較
• Redis モジュールの適用範囲
• Azure Cache for Redis の価格
  
  ※執筆時点

3.2. ベクトルデータの展開

3.2.1 Redis アクセスキーの取得

Azure Portal からアクセスキー (プライマリかセカンダリのどちらか) をコピーします。キーは絶対に外部に漏らさないように注意します。

3.2.2. 環境変数の設定

Azure Cache for Redis リソース名とアクセスキーを事前に環境変数に登録しておきます。

# For Ubuntu
# !export REDIS_NAME=<your-redis-name>
# !export REDIS_KEY=<your-redis-key>

# For Windows
# !set REDIS_NAME=<your-redis-name>
# !set REDIS_KEY=<your-redis-key>

3.2.2. インデックス作成

下記部分のコードを実行すると Redis への接続とベクトル検索のためのインデックス作成を行います。Approximate Nearest Neighbor (ANN または近似最近傍探索) を行うために、Hierarchical Navigable Small World (HNSW) インデックスを選択しています。今回行っているテキスト検索ではではつまるところテキストから生成したベクトルどうしの最近傍探索を行って互いに近いテキストを探しています。近似最近傍探索では、厳密な最近傍ではないものの高速に探索を行うことができます。おそらく、Redis を使う必要があるほどのデータサイズを扱う状況では検索時間を早めたい場合の方が多いのではないかと思います。もしも、厳密な最近傍探索を行いたい場合は総当たりの FLAT インデックスを選択することもできます。

# redis_name = "<your-redis-name>"
redis_name = os.environ["REDIS_NAME"]
redis_host = f"{redis_name}.southcentralus.redisenterprise.cache.azure.net"  # Redis が米国中南部リージョンに作成されている場合の例
# redis_key =  "<your-redis-key>"
redis_key = os.environ["REDIS_KEY"]

# Redis へ接続
redis_conn = redis.StrictRedis(host=redis_host,port=10000, password=redis_key, ssl=True)

# インデックスを作成する
schema = ([
    VectorField("Embedding", "HNSW", {"TYPE": "FLOAT32", "DIM": embedding_dimension, "DISTANCE_METRIC": "COSINE"}),  # RediSearch は cosine DISTANCE を使用
    TextField("ProductId"),
    TextField("UserId"),
    NumericField("Score"),
    TextField("Summary"),
    TextField("Text"),
    TextField("Combined"),
    NumericField("N_tokens")
])
# redis_conn.ft().dropindex(schema)   # インデックスをdorp-createする場合はコメントを除く
redis_conn.ft().create_index(schema)

参考

• RedisDays New York 2022: Using AI to Reveal Trading Signals Buried in Corporate Filings (YouTube)
• Vector similarity (document)
• セカンダリインデックス
• redis-py - Python Client for Redis
• RedisAI/vecsim-demo/SemanticSearch1k.ipynb
• Vector Similarity (sample)
• Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs

3.2.3. ベクトルデータの展開

インデックスを作成したフィールドにデータを展開します。生成したベクトルデータと共に、元となったテキストなどいくつかの情報を合わせて展開します。

# Redis に展開
for i, row in df.iterrows():
    d = {
        "Embedding": np.array(row["Embedding"]).astype(np.float32).tobytes(),
        "ProductId": row["ProductId"],
        "UserId":    row["UserId"],
        "Score":     row["Score"],
        "Summary":   row["Summary"],
        "Text":      row["Text"],
        "Combined":  row["Combined"],
        "N_tokens":  row["N_tokens"]
    }
    redis_conn.hset(str(i), mapping=d)

4. テキスト検索 (On Redis)

4.1. 関数定義

検索を実行する関数を定義します。まず、ローカル PC 上での実行時と時と同じく検索クエリのテキストを与えると Azure OpenAI Service API を呼んでクエリのベクトル化を行ます。続けて、RediSearch が近似最近傍探索を行ってクエリと近い内容の商品レビューをデフォルトで 3 件返します。
なお、RediSearch では 3 種類のベクトル間距離メトリックをサポートしています。今回はそのうちのひとつであるコサイン距離を選択して、結果をコサイン類似度に変換して返しています。

# クエリにマッチする商品を検索して返す (Redis 上)
def search_reviews_redis(query, n=3, pprint=True, engine="text-search-ada-query-001"):
    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("; Content:", ": ")
        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

クエリ書式について

Vector similarity クエリはクセが強いです。色々なクエリを試したい場合は下記ドキュメントを参考にします。

• Vector search examples

4.2. 検索実行 (On Redis)

"delicious beans (おいしいお豆)" で検索します。

results_redis = search_reviews_redis("delicious beans", n=3, engine=embedding_model_for_query)

ローカル PC 上での検索と同じフォーマットで結果が返ってきます。

0.375126361847 | Delicious!:  I enjoy this white beans seasoning, it gives a rich flavor to the beans I just love it, my mother in law didn't know about this Zatarain's brand and now she is traying diff

0.37370193004600005 | Jamaican Blue beans:  Excellent coffee bean for roasting. Our family just purchased another 5 pounds for more roasting. Plenty of flavor and mild on acidity when roasted to a dark brown

0.373187303543 | Good Buy:  I liked the beans. They were vacuum sealed, plump and moist. Would recommend them for any use. I personally split and stuck them in some vodka to make vanilla extract. Yum!

ローカル PC 上での検索結果と比較してみます。

results = search_reviews(df, "delicious beans", n=3, engine=embedding_model_for_query)

近似最近傍探索のため全く同じ結果にはなりませんが、おおむね同じ結果を返せていることがわかります。

0.3770173260434461 | Best beans your money can buy:  These are, hands down, the best jelly beans on the market.  There isn't a gross one in the bunch and each of them has an intense, delicious flavor.  Though I hesitate t

0.37512639635782036 | Delicious!:  I enjoy this white beans seasoning, it gives a rich flavor to the beans I just love it, my mother in law didn't know about this Zatarain's brand and now she is traying different seasoning

0.37370195283798296 | Jamaican Blue beans:  Excellent coffee bean for roasting. Our family just purchased another 5 pounds for more roasting. Plenty of flavor and mild on acidity when roasted to a dark brown bean and befor

おわりに

今回は Azure OpenAI Service のモデルの中でも比較的パターン化しやすい Embeddings モデルに着目したシナリオについて記事を書きました。OpenAI や Azure OpenAI Service の登場により、これまでは高い専門知識・技術を要していた超巨大言語モデルを扱うための障壁が大きく下がり、これらのモデルになじみのなかった利用者層も手軽に触れることができるようになりました。生成系モデルも含むこうしたモデルの活用は利用者の独創性次第で非常に大きな可能性を秘めていると思います。今後、革新的な活用方法が数多く生み出されることを想像すると楽しみです。

以上です。🍵