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

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

今回は私が好きな [Azure Cogntive Search](https://www.youtube.com/watch?v=9GcjFofXdd0) と Azure OpenAI を組み合わせるといったいどのようなことができるようになるのか、いろいろ試してみたいと思います。

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

https://qiita.com/nohanaga/items/f710cac82072b63bc73f


# 組み合わせ方法
Azure Cogntive Search をメインにして[カスタムスキル](https://qiita.com/nohanaga/items/986887b82cf76dc7b0a9)を使えば Azure OpenAI の様々な機能を統合することができます。

![005.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/674344/ac7680de-3e78-2a2d-beae-29ac57e8f5f1.png)

例えば、長いドキュメントをチャンクに分割し、要約し、検索結果に出力する。要約の仕方もプロンプトエンジニアリングによって、[さまざまな形式](https://learn.microsoft.com/azure/cognitive-services/openai/how-to/completions#summarization)で出力させることも可能になります。他にも、

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

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

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

# Embeddings API
Azure OpenAI ではテキストの[埋め込み](https://learn.microsoft.com/azure/cognitive-services/openai/how-to/embeddings?tabs=console)（Embeddings）を簡単に生成できます。このベクトルを使うことによってテキスト間の類似性を計算することができます。類似性の計算にはコサイン類似度を使います。

https://learn.microsoft.com/ja-jp/azure/cognitive-services/openai/tutorials/embeddings

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

# ランカーアーキテクチャ

ここで Azure Cognitive Search の[ランカー](https://learn.microsoft.com/azure/search/semantic-search-overview#order-of-operations)と、Azure OpenAI によって作るランカーの違いを見ておきましょう。

![001.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/674344/0fec14db-42a8-8c3c-5d6e-8d9d780ce6d5.png)

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

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

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

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

# 手順（まずは検証）
1. Azure OpenAI アカウントの作成（2023/02 現在 [申請制]()）本家の OpenAI でも OK
1. Python 環境の用意、必要ライブラリのインストール
    - `openai` ライブラリは本家も Azure も共通です
    - `transformers` ライブラリは[トークンの計算](https://platform.openai.com/tokenizer)に利用します
※[こちら](https://learn.microsoft.com/azure/cognitive-services/openai/tutorials/embeddings)のチュートリアルの手順通りに行えば、Embeddings の生成まで行えます
1. [Azure Cache for Redis Enterprise](https://learn.microsoft.com/azure/azure-cache-for-redis/quickstart-create-redis-enterprise) の作成（RediSearch モジュールにチェック）
1. 検索するデータセットの準備と埋め込み生成
以前 Recommender Workshop で使用した [MovieLens](https://github.com/nohanaga/movie_recommendation_cos_similarity/blob/main/movie_dataset.csv) Dataset(4803 movies) を使います。
1. Redis インデックス構築
1. ベクトル検索

## 1. 検索するデータセットの準備と埋め込み生成
[こちら](https://learn.microsoft.com/azure/cognitive-services/openai/tutorials/embeddings)のチュートリアルを参考にベクトル化する Dataframe のカラムを結合して今回は `title`,  `keywords`,  `genres`, `cast` カラムを結合して `combined_features` カラムを作成しました。

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

```python
%%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` を取得します。
```python
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 つのインデックス方法を[サポート](https://redis.io/docs/stack/search/reference/vectors/)しています。
- FLAT: ブルートフォース インデックス
- HNSW: Hierarchical Navigable Small World(HNSW) を使用した効率的で堅牢な近似最近傍検索の実装です。

アルゴリズムの解説や使い分けについては[こちら](https://yusukematsui.me/project/survey_pq/doc/ann_billion_2018.pdf)の資料が最高でした。N=100 万程度までは FLAT がおすすめとのこと。また、`DISTANCE_METRIC` として、`L2`、`IP`、`COSINE` を選択可能です。今回の場合 `COSINE` を使います。

```python
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](https://redis.io/docs/stack/search/reference/vectors/#creation-attributes-per-algorithm) 上では ベクトルタイプが `FLOAT64` にも対応していると記載がありましたが、実際に指定するとエラーが出て作成できませんでした。そのため、`FLOAT32` にしています。

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

```python
%%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 クエリーはぱっと見分かりにくいので[サンプル](https://redis.io/docs/stack/search/reference/vectors/#vector-search-examples)を参考にしてください。

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

```python
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 に近いほどベクトルが類似していることを表します。

```python
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` に含まれるので当然か。

```python
search_docs("元軍人の父が娘を救出するため、単身で島に上陸する", 3)
```
2996 0.724949836731 Commando
2521 0.742291748524 Tidal Wave
1404 0.760077595711 Street Fighter

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

```python
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` に含まれない。日本語で「起きろ、ネオ」はダメでした。

```python
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` に含まれない。


```python
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 にデプロイして東日本からアクセスしたので地理的な遅延が加算されます。

#### 埋め込み生成
```python
%%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)

#### クエリー実行
```python
%%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 種類のモデルがあり、精度と速度のトレードオフとなっています。精度を犠牲に高速にするには[ほかのモデル](https://learn.microsoft.com/azure/cognitive-services/openai/concepts/models#embeddings-models-1)を使うこともできます。今回は一番精度の高い Davinci を使用しています。

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

# ハイブリッド検索
当然頻度ベースの検索とベクトル検索を同時に行いたいというニーズが出てくるわけですが、現状 Redis では日本語の全文検索ができません。（[フィルター](https://redis.io/docs/stack/search/reference/vectors/#hybrid-knn-queries)は使えます）ぐぬぬ。。。これを同時に行うことができる検索エンジンとなると、現状 Apache Lucene と ElasticSearch くらいしかないのでは？

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

# 参考
https://github.com/RediSearch/RediSearch/blob/master/docs/docs/vecsim-range_queries_examples.ipynb

https://github.com/RedisAI/vecsim-demo/blob/master/SemanticSearch1k.ipynb