Azure OpenAI と Azure Cognitive Search の組み合わせを考える

はじめに

巷では OpenAI が人気ですね。OpenAI の GPT-3 モデルを Azure 上でホスティングしたサービスである Azure OpenAI Service の GA によって企業はエンタープライズ用途で GPT-3 モデルを使用することができるようになりました。これにより Azure Cognitive Services と同様のセキュリティ機能、SLA、コンプライアンス、責任ある AI などのメリットを享受できます。

今回は私が好きな Azure Cogntive Search と Azure OpenAI を組み合わせるといったいどのようなことができるようになるのか、いろいろ試してみたいと思います。

🎉5/24 ついにベクトル検索機能が実装されました！

組み合わせ方法

Azure Cogntive Search をメインにしてカスタムスキルを使えば Azure OpenAI の様々な機能を統合することができます。

例えば、長いドキュメントをチャンクに分割し、要約し、検索結果に出力する。要約の仕方もプロンプトエンジニアリングによって、さまざまな形式で出力させることも可能になります。他にも、

• 文章分類
• キーフレーズ抽出
• 感情分析
• 固有名詞抽出
• 翻訳

などのタスクもこなしてしまいますが、もともと組み込まれている Text Analytics スキルとかちあっちゃいますね…　この点はユーザーのデータや求める精度、予算などに基づいて使い分ける必要があります。ぜひ検証をお願いいたします。

そして、検索ランキングを生成する部分も Azure OpenAI によって独自に用意することが可能になります。今日はこちらをメインに話していこうと思います。

Embeddings API

Azure OpenAI ではテキストの埋め込み（Embeddings）を簡単に生成できます。このベクトルを使うことによってテキスト間の類似性を計算することができます。類似性の計算にはコサイン類似度を使います。

こちらのチュートリアルを使えば、実験室環境でドキュメント検索を行うことができます。

ランカーアーキテクチャ

ここで Azure Cognitive Search のランカーと、Azure OpenAI によって作るランカーの違いを見ておきましょう。

Azure Cognitive Search は Apache Lucene をベースとして開発された全文検索エンジンで、スコアリングには BM25Similarity の TF-IDF タイプの頻度ベースのランカーが実装されています。ヒットさせるには検索キーワードがドキュメントに含まれていなければなりません。次のランカーとして「セマンティック検索(Preview)」という Microsoft Research が開発している Turing ベースの大規模言語モデル (LLM) がリランカーとして実装されています。リランカーはあくまで一致したキーワードの上位 50 件を並び替えるという目的で実装されています。この構造のため、ユーザーはドキュメント内に含まれているキーワードで検索する必要があります。検索キーワードには含まれていない、ユーザーの意図と近いワードをヒットさせることはできません。

うーん、キーワード一致を使いつつ、GPU を大量に使う LLM の検索リソース削減、それでいて既存の検索結果の精度をさらに上げるという絶妙なバランスの上でこのように実装されてるんだなぁ…　とか考えさせられます。今後、いい感じにキーワード一致とセマンティック検索を組み合わせてみたい。

まぁ今回はせっかく Azure OpenAI で Embeddings API が使えるということで、最初のランカーとして GPT-3 によるベクトル検索を on Azure でやってみようと思います。ですが、Azure Cognitive Search では 2023/02 現在、ベクトル検索の機能が実装されていません😭 最近、Apache Lucene では ベクトルの近似最近傍探索(ANN Search) が実装されたと聞きましたが、これが Azure Cognitive Search に実装されるにはさらに時間がかかることでしょう・・・

ということで、 on Azure でベクトル検索を行う方法として、一つは Azure Cache for Redis を使用するという案。Azure 上でのフルマネージドインメモリ DB、ナイスですね！ただし、ベクトル検索を行うには、RediSearch というモジュールが必要となりこれは Enterprise プランでしか利用できません。Enterprise プランの価格、高いです。個人の場合、VM を立てて Redis Stack Server を構築することで検証可能です。

手順（まずは検証）

1. Azure OpenAI アカウントの作成（2023/02 現在 申請制）本家の OpenAI でも OK
2. Python 環境の用意、必要ライブラリのインストール
   • openai ライブラリは本家も Azure も共通です
   • transformers ライブラリはトークンの計算に利用します
     ※こちらのチュートリアルの手順通りに行えば、Embeddings の生成まで行えます
3. Azure Cache for Redis Enterprise の作成（RediSearch モジュールにチェック）
4. 検索するデータセットの準備と埋め込み生成
   以前 Recommender Workshop で使用した MovieLens Dataset(4803 movies) を使います。
5. Redis インデックス構築
6. ベクトル検索

1. 検索するデータセットの準備と埋め込み生成

こちらのチュートリアルを参考にベクトル化する Dataframe のカラムを結合して今回は title, keywords, genres, cast カラムを結合して combined_features カラムを作成しました。

埋め込みは、get_embedding を使用します。

%%time
df['davinci_search'] = df["combined_features"].apply(lambda x : get_embedding(x, engine = 'text-search-davinci-doc-001'))

CPU times: user 19 s, sys: 2.61 s, total: 21.7 s
Wall time: 22min 33s

埋め込みの生成に text-search-davinci-doc-001 を使っているので速度は遅いです。
毎度 get_embedding したくないので、Dataframe を Parquet などで保存しておきます。

2. Redis ベクトルインデックス構築

2.1. Redis への接続

Azure Portal から Redis Enterprise リソースを開き、Endpoint と Access Key を取得します。

from redis import Redis
from redis.commands.search.field import VectorField, TagField, NumericField, TextField
from redis.commands.search.query import Query
from redis.commands.search.result import Result

redis_conn = redis.StrictRedis(host='<Your Endpoint>',
    port=10000, db=0, password='<Your Key>', ssl=True)

redis_conn.ping()

2.2. ベクトルインデックスの作成

RediSearch は、次の 2 つのインデックス方法をサポートしています。

• FLAT: ブルートフォース インデックス
• HNSW: Hierarchical Navigable Small World(HNSW) を使用した効率的で堅牢な近似最近傍検索の実装です。

アルゴリズムの解説や使い分けについてはこちらの資料が最高でした。N=100 万程度までは FLAT がおすすめとのこと。また、DISTANCE_METRIC として、L2、IP、COSINE を選択可能です。今回の場合 COSINE を使います。

def create_flat_index (redis_conn,vector_field_name,number_of_vectors, vector_dimensions=12288, distance_metric='COSINE'):
    redis_conn.ft().create_index([
        VectorField(vector_field_name, "FLAT", {"TYPE": "FLOAT32", "DIM": vector_dimensions, "DISTANCE_METRIC": distance_metric, "INITIAL_CAP": number_of_vectors, "BLOCK_SIZE":number_of_vectors }),
        TagField("itemid"),
        TextField("combined_features"),
        TextField("title")
    ])
    
def create_hnsw_index (redis_conn,vector_field_name, number_of_vectors, vector_dimensions=12288, distance_metric='COSINE',M=40,EF=200):
    redis_conn.ft().create_index([
        VectorField(vector_field_name, "HNSW", {"TYPE": "FLOAT32", "DIM": vector_dimensions, "DISTANCE_METRIC": distance_metric, "INITIAL_CAP": number_of_vectors, "M": M, "EF_CONSTRUCTION": EF}),
        TagField("itemid"),
        TextField("combined_features"),
        TextField("title")
    ])    


ITEM_KEYWORD_EMBEDDING_FIELD='davinci_search'
#Davinci is 12288 dims
TEXT_EMBEDDING_DIMENSION=12288
NUMBER_PRODUCTS=4803

# create flat index & load vectors
create_flat_index(redis_conn, ITEM_KEYWORD_EMBEDDING_FIELD, NUMBER_PRODUCTS, TEXT_EMBEDDING_DIMENSION, 'COSINE')

Redis の Docs 上では ベクトルタイプが FLOAT64 にも対応していると記載がありましたが、実際に指定するとエラーが出て作成できませんでした。そのため、FLOAT32 にしています。

2.3. ベクトルをハッシュに格納

検索結果の内容を分析するために、combined_features カラムも登録しています。登録時、ベクトルは Byte 列に変換する必要があります。

%%time
def load_vectors(client:Redis, df, vector_field_name):
    p = client.pipeline(transaction=False)
    
    for i in df.index:
        key= str(i)
        item_keywords_vector = df[vector_field_name][i].astype(np.float32).tobytes()
        item_metadata = {'itemid': str(df['id'][i]), 'title': df['title'][i], 'combined_features': df['combined_features'][i], vector_field_name: item_keywords_vector}
        
        client.hset(key, mapping=item_metadata)

    p.execute()

load_vectors(redis_conn, moviedf, ITEM_KEYWORD_EMBEDDING_FIELD)

Wall time: 13min 44s

2.4. ベクトル検索(ANN Search)

検索キーワードの埋め込みを得るには get_embedding でクエリー用のモデル text-search-davinci-query-001 を指定します。Redis の kNN クエリーはぱっと見分かりにくいのでサンプルを参考にしてください。

いくつかの有名映画をヒットさせるための検索クエリーをあらすじから抜粋して記述します。

def search_docs(user_query, topK):
    embedding = get_embedding(user_query, engine="text-search-davinci-query-001")
    query_byte_array = np.array(embedding).astype(np.float32).tobytes()

    q = Query(f'*=>[KNN {topK} @{ITEM_KEYWORD_EMBEDDING_FIELD} $vec AS vector_score]').return_fields("vector_score", "title").sort_by('vector_score').paging(0,topK).dialect(2)
    results = redis_conn.ft().search(q, query_params={"vec":  query_byte_array})

    for item in results.docs:
        print (item.id, item.vector_score, item.title)

search_docs("AI and humanity will confront each other", 5)

634 0.734805703163 The Matrix
1997 0.738143444061 Her
266 0.738558232784 I, Robot
43 0.739533424377 Terminator Salvation
2654 0.741096973419 Automata

「AI」「humanity」という単語は combined_features には含まれていませんが、いい感じでランキング生成されています！

※ Redis の COSINE は コサイン距離であり、コサイン類似度を 1 から引いた値になっています。つまり 0 に近いほどベクトルが類似していることを表します。

search_docs("Ex-military father lands on island alone to rescue daughter", 3)

2996 0.688981413841 Commando
1331 0.714794993401 Nim's Island
2953 0.715987503529 Return to the Blue Lagoon

面白い結果だな、気に入った。「daughter」「father」「rescue」 が combined_features に含まれるので当然か。

search_docs("元軍人の父が娘を救出するため、単身で島に上陸する", 3)

2996 0.724949836731 Commando
2521 0.742291748524 Tidal Wave
1404 0.760077595711 Street Fighter

日本語でもポンッとできました。トリックではないです、日本語はデータセットに含まれません。

search_docs("Wake up, Neo", 5)

634 0.718291580677 The Matrix
123 0.729486823082 The Matrix Revolutions
125 0.731914699078 The Matrix Reloaded
1153 0.756486654282 Lucy
3207 0.757386445999 Awake

「Neo」は combined_features に含まれない。日本語で「起きろ、ネオ」はダメでした。

search_docs("Prof. Langdon investigates the mystery of Louvre.", 5)

201 0.718103647232 The Da Vinci Code
3082 0.718517065048 Sphinx
1449 0.727268278599 The Order
128 0.731719851494 Angels & Demons
291 0.736838400364 National Treasure

「Langdon」も「Louvre」も combined_features に含まれない。

search_docs("ラングドン教授がルーブル美術館の謎を探る", 5)

291 0.775311648846 National Treasure
201 0.776099085808 The Da Vinci Code
483 0.778588891029 Timeline
3615 0.779091000557 The Barbarian Invasions
2389 0.779899418354 Renaissance

2.5. 実行時間

もう少し細かく時間を測ります。OpenAI リソースは South Central US に、Redis は East US にデプロイして東日本からアクセスしたので地理的な遅延が加算されます。

埋め込み生成

%%timeit
user_query = "AI and humanity will confront each other"
embedding = get_embedding(user_query, engine="text-search-davinci-query-001")

259 ms ± 40.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

クエリー実行

%%timeit
topK=5
query_byte_array = np.array(embedding).astype(np.float32).tobytes()

q = Query(f'*=>[KNN {topK} @{ITEM_KEYWORD_EMBEDDING_FIELD} $vec AS vector_score]').return_fields("vector_score", "title").sort_by('vector_score').paging(0,topK).dialect(2)
results = redis_conn.ft().search(q, query_params={"vec":  query_byte_array})

for item in results.docs:
    print (item.id, item.vector_score, item.title)

182 ms ± 703 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

埋め込みの生成が時間かかりますね。GPT-3 モデルには大きく分けて 4 種類のモデルがあり、精度と速度のトレードオフとなっています。精度を犠牲に高速にするにはほかのモデルを使うこともできます。今回は一番精度の高い Davinci を使用しています。

また、大規模データセットの場合は Spark を使うなどして分散処理を行うことでさらに高速化は可能です。（API limit内で）

ハイブリッド検索

当然頻度ベースの検索とベクトル検索を同時に行いたいというニーズが出てくるわけですが、現状 Redis では日本語の全文検索ができません。（フィルターは使えます）ぐぬぬ。。。これを同時に行うことができる検索エンジンとなると、現状 Apache Lucene と ElasticSearch くらいしかないのでは？

さいごに

今回は Redis でベクトル検索をするところまでとさせていただきますが、次回以降も連携方法を考えていきたいと思います。Azure Cognitive Search を使っている方でベクトル検索にご興味がある方、「ベクトル検索実装はよ！」 運動にご協力をお願い致します。

参考