MongoDB AtlasのVector Searchを使用してみた

はじめに

MongoDB Atlasとベクトルデータベース（Pinecone）を併用して開発していたところ、MongoDB AtlasのVector Search機能をの機能を使用することで、データベースの構成をMongoDB Atlasのみに簡素化できました。
Vector Searchの追加設定と、Pythonのコードについて記載します。

Doucmentにembedding vectorを登録

既存のMongodb AtlasのDocuemntに新たに画像のembedding vectorを追加しました。
Collectionからとりだしたdocumentにembedding vectorを追加したのち、もとのdoucmentに上書きします。
追加した項目名（vector_database_field_name)は"embed"にしました。

            vector_database_field_name="embed"
            imageFeature_np = get_single_image_embedding(Image.open(filename))
            imageEmbedding = imageFeature_np[0].tolist()

            if vector_database_field_name not in webCamInfo:
                webCamInfo[vector_database_field_name] = imageEmbedding
            
            collection.replace_one({'_id': webCamInfo['_id']}, webCamInfo)

画像のembedding化には、「openai/clip-vit-base-patch32」のモデルを使用しました。

from transformers import AutoProcessor,CLIPProcessor, CLIPModel, CLIPTokenizer

model_ID = "openai/clip-vit-base-patch32"

def get_model_info(model_ID, device):
	model = CLIPModel.from_pretrained(model_ID).to(device)
	processor = AutoProcessor.from_pretrained(model_ID)
	tokenizer = CLIPTokenizer.from_pretrained(model_ID)

model, processor, tokenizer = get_model_info(model_ID, device)

def get_single_image_embedding(my_image):
    image = processor(images=my_image , return_tensors="pt")
    embedding = model.get_image_features(**image).float()
    # convert the embeddings to numpy array
    return embedding.cpu().detach().numpy()

Indexの作成

Mongodb Atlasのコンソールの「Atlas Search」から、追加した"embed"にたいしてIndexを作成することができいます。Json形式で、Indexの仕様を定義できます。
「openai/clip-vit-base-patch32」のdimensionは512、検索時には、Cosine Similarityを使用するように定義しました。

{
  "fields": [
    {
      "numDimensions": 512,
      "path": "embed",
      "similarity": "cosine",
      "type": "vector"
    }
  ]
}

作成した結果、MongoDB Atlasのコンソールは以下のように表示されます。Indexの名前は、「imageindex」にしています。

検索について

embedding vectorを検索する際は、MongoDB の aggregate 関数と $vectorSearch 演算子を使用します。

    vector_database_field_name="embed"
    imageFeature_np = get_single_text_embedding(text)
    imageEmbedding = imageFeature_np[0].tolist()
    results = collection.aggregate([
        {
            "$vectorSearch": {
                "queryVector": imageEmbedding,
                "path": vector_database_field_name,
                "numCandidates": 100, # this should be 10-20x the limit
                "limit": 10, 
                "index": 'imageindex', # the index name you used in the earlier step
            }
        }
    ])

検索の主要な要素

1. Index名: imageindex
   作成したIndexの名前を指定します。

2. queryVector: 検索ベクトル
   これは検索に使用するベクトル。

3. path: "embed"
   これは、ドキュメント内でベクトルデータが格納されているフィールド名。"embed" というフィールド名を使用しています。

4. limit: 10
   検索する件数

5. numCandidates: 100
   候補ベクトルの最大数。

また、検索した結果をさらにフィルタリングする等、演算子と使用することができます。ただし、$vectorSearchは先頭でなければならないようです。

    results = collection.aggregate([
        {
            "$vectorSearch": {
                "queryVector": imageEmbedding,
                "path": vector_database_field_name,
                "numCandidates": 100, # this should be 10-20x the limit
                "limit": 10, # the number of documents to return in the results
                "index": 'imageindex', # the index name you used in the earlier step
            }
        }
        {
            "$match":{
                "location.country": 'United States'
            }
        }
    ])

おわりに

MongoDB Atlasとベクトルデータベースを併用していたところ、Vector Searchを導入することで、データベースの構成だけでなく、プログラムのほうも簡潔になりました。さらに、aggregate関数を使用することで、より複雑な検索もできそうです。今後も、うまく利用していきたいと思いました。