Vector Retrieval Service for Milvus:Implementasikan pencarian multimodal dengan Milvus dan Qwen

Informasi latar belakang

Dalam pencarian multimodal, data tidak terstruktur seperti gambar dan teks diubah menjadi representasi vektor, lalu teknologi pencarian vektor digunakan untuk menemukan konten yang mirip. Topik ini menggunakan alat-alat berikut:

• Milvus: Database vektor yang efisien untuk menyimpan dan mengambil vektor.

• Qwen-VL: Mengekstraksi deskripsi dan kata kunci dari gambar. Untuk informasi lebih lanjut, lihat Qwen-VL.

• DashScope Embedding API: Mengonversi gambar dan teks menjadi vektor. Untuk informasi lebih lanjut, lihat Detail API Multimodal-Embedding.

Mode pencarian yang didukung meliputi:

• Pencarian teks-ke-gambar: Masukkan kueri teks untuk menemukan gambar yang paling mirip.

• Pencarian teks-ke-teks: Masukkan kueri teks untuk menemukan deskripsi gambar yang paling mirip.

• Pencarian berdasarkan gambar: Masukkan kueri gambar untuk menemukan gambar yang paling mirip.

• Pencarian gambar-ke-teks: Masukkan kueri gambar untuk menemukan deskripsi gambar yang paling mirip.

Arsitektur sistem

Gambar berikut menunjukkan arsitektur keseluruhan sistem pencarian multimodal.

Prasyarat

• Anda telah membuat instans Milvus. Untuk informasi lebih lanjut, lihat Buat instans Milvus.

• Anda telah mengaktifkan Alibaba Cloud Model Studio dan memperoleh Kunci API. Untuk informasi lebih lanjut, lihat Peroleh Kunci API.

• Anda telah menginstal paket dependensi yang diperlukan.

  pip3 install dashscope pymilvus==2.5.0

  Contoh dalam topik ini dijalankan di lingkungan Python 3.9.

• Anda telah mengunduh dan mengekstrak set data sampel.

  wget https://github.com/milvus-io/pymilvus-assets/releases/download/imagedata/reverse_image_search.zip
  unzip -q -o reverse_image_search.zip

  Set data sampel berisi file CSV bernama reverse_image_search.csv dan beberapa file gambar.

  Catatan

  Set data sampel dan gambarnya yang digunakan dalam topik ini berasal dari proyek open source Milvus.

Pengenalan kode inti

Dalam contoh ini, model Qwen-VL mengekstraksi deskripsi gambar dan menyimpannya di bidang image_description. Model penyematan multimodal kemudian mengonversi gambar dan deskripsinya menjadi representasi vektor, seperti image_embedding dan text_embedding, untuk mempersiapkan data guna pencarian lintas-modalitas.

Untuk menyederhanakan demo, hanya 200 gambar pertama yang digunakan.

import base64
import csv
import dashscope
import os
import pandas as pd
import sys
import time
from tqdm import tqdm
from pymilvus import (
    connections,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    MilvusException,
    utility,
)

from http import HTTPStatus
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class FeatureExtractor:
    def __init__(self, DASHSCOPE_API_KEY):
        self._api_key = DASHSCOPE_API_KEY  # Gunakan variabel lingkungan untuk menyimpan kunci API

    def __call__(self, input_data, input_type):
        if input_type not in ("image", "text"):
            raise ValueError("Invalid input type. Must be 'image' or 'text'.")

        try:
            if input_type == "image":
                _, ext = os.path.splitext(input_data)
                image_format = ext.lstrip(".").lower()
                with open(input_data, "rb") as image_file:
                    base64_image = base64.b64encode(image_file.read()).decode("utf-8")
                input_data = f"data:image/{image_format};base64,{base64_image}"
                payload = [{"image": input_data}]
            else:
                payload = [{"text": input_data}]

            resp = dashscope.MultiModalEmbedding.call(
                model="multimodal-embedding-v1",
                input=payload,
                api_key=self._api_key,
            )

            if resp.status_code == HTTPStatus.OK:
                return resp.output["embeddings"][0]["embedding"]
            else:
                raise RuntimeError(
                    f"API call failed. Status code: {resp.status_code}, Error message: {resp.message}"
                )
        except Exception as e:
            logger.error(f"Processing failed: {str(e)}")
            raise


class FeatureExtractorVL:
    def __init__(self, DASHSCOPE_API_KEY):
        self._api_key = DASHSCOPE_API_KEY  # Gunakan variabel lingkungan untuk menyimpan kunci API

    def __call__(self, input_data, input_type):
        if input_type not in ("image"):
            raise ValueError("Invalid input type. Must be 'image'.")

        try:
            if input_type == "image":
                payload=[
                            {
                                "role": "system",
                                "content": [{"type":"text","text": "You are a helpful assistant."}]
                            },
                            {
                                "role": "user",
                                "content": [
                                            # {"image": "https://dashscope.oss-cn-beijing.aliyuncs.com/images/dog_and_girl.jpeg"},
                                            {"image": input_data},
                                            {"text": "First, describe this image in under 50 words, and then provide 5 keywords"}
                                            ],
                            }
                        ]

            resp = dashscope.MultiModalConversation.call(
                model="qwen-vl-plus",
                messages=payload,
                api_key=self._api_key,
            )

            if resp.status_code == HTTPStatus.OK:
                return resp.output["choices"][0]["message"].content[0]["text"]
            else:
                raise RuntimeError(
                    f"API call failed. Status code: {resp.status_code}, Error message: {resp.message}"
                )
        except Exception as e:
            logger.error(f"Processing failed: {str(e)}")
            raise


class MilvusClient:
    def __init__(self, MILVUS_TOKEN, MILVUS_HOST, MILVUS_PORT, INDEX, COLLECTION_NAME):
        self._token = MILVUS_TOKEN
        self._host = MILVUS_HOST
        self._port = MILVUS_PORT
        self._index = INDEX
        self._collection_name = COLLECTION_NAME

        self._connect()
        self._create_collection_if_not_exists()

    def _connect(self):
        try:
            connections.connect(alias="default", host=self._host, port=self._port, token=self._token)
            logger.info("Connected to Milvus successfully.")
        except Exception as e:
            logger.error(f"Failed to connect to Milvus: {str(e)}")
            sys.exit(1)

    def _collection_exists(self):
        return self._collection_name in utility.list_collections()
    
    def _create_collection_if_not_exists(self):
        try:
            if not self._collection_exists():
                fields = [
                    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
                    FieldSchema(name="origin", dtype=DataType.VARCHAR, max_length=512),
                    FieldSchema(name="image_description", dtype=DataType.VARCHAR, max_length=1024),
                    FieldSchema(name="image_embedding", dtype=DataType.FLOAT_VECTOR, dim=1024),
                    FieldSchema(name="text_embedding", dtype=DataType.FLOAT_VECTOR, dim=1024)
                ]

                schema = CollectionSchema(fields)

                self._collection = Collection(self._collection_name, schema)

                if self._index == 'IVF_FLAT':
                    self._create_ivf_index()
                else:
                    self._create_hnsw_index()   
                logger.info("Collection created successfully.")
            else:
                self._collection = Collection(self._collection_name)
                logger.info("Collection already exists.")
        except Exception as e:
            logger.error(f"Failed to create or load the collection: {str(e)}")
            sys.exit(1)


    def _create_ivf_index(self):
        index_params = {
            "index_type": "IVF_FLAT",
            "params": {
                        "nlist": 1024, # Jumlah kluster untuk indeks
                    },
            "metric_type": "L2",
        }
        self._collection.create_index("image_embedding", index_params)
        self._collection.create_index("text_embedding", index_params)
        logger.info("Index created successfully.")

    def _create_hnsw_index(self):
        index_params = {
            "index_type": "HNSW",
            "params": {
                        "M": 64, # Jumlah maksimum tetangga yang dapat dihubungkan setiap node dalam graf
                        "efConstruction": 100, # Jumlah tetangga kandidat yang dipertimbangkan saat pembuatan indeks
                    },
            "metric_type": "L2",
        }
        self._collection.create_index("image_embedding", index_params)
        self._collection.create_index("text_embedding", index_params)
        logger.info("Index created successfully.")
    
    def insert(self, data):
        try:
            self._collection.insert(data)
            self._collection.load()
            logger.info("Data inserted and loaded successfully.")
        except MilvusException as e:
            logger.error(f"Failed to insert data: {str(e)}")
            raise

    def search(self, query_embedding, field, limit=3):
        try:
            if self._index == 'IVF_FLAT':
                param={"metric_type": "L2", "params": {"nprobe": 10}}
            else:
                param={"metric_type": "L2", "params": {"ef": 10}}

            result = self._collection.search(
                data=[query_embedding],
                anns_field=field,
                param=param,
                limit=limit,
                output_fields=["origin", "image_description"],
            )
            return [{"id": hit.id, "distance": hit.distance, "origin": hit.origin, "image_description": hit.image_description} for hit in result[0]]
        except Exception as e:
            logger.error(f"Search failed: {str(e)}")
            return None


# Muat data dan hasilkan embedding
def load_image_embeddings(extractor, extractorVL, csv_path):
    df = pd.read_csv(csv_path)
    image_embeddings = {}

    for image_path in tqdm(df["path"].tolist()[:200], desc="Generating image embeddings"): # Gunakan hanya 200 gambar pertama untuk demo
        try:
            desc = extractorVL(image_path, "image")
            image_embeddings[image_path] = [desc, extractor(image_path, "image"), extractor(desc, "text")]
            time.sleep(1)  # Kendalikan frekuensi pemanggilan API
        except Exception as e:
            logger.warning(f"Failed to process {image_path}, skipping: {str(e)}")

    return [{"origin": k, 'image_description':v[0], "image_embedding": v[1], 'text_embedding': v[2]} for k, v in image_embeddings.items()]
    

Di mana:

• FeatureExtractor: Memanggil DashScope Embedding API untuk mengonversi gambar atau teks menjadi representasi vektor.

• FeatureExtractorVL: Memanggil model Qwen-VL untuk mengekstraksi deskripsi teks dan kata kunci dari gambar.

• MilvusClient: Mengenkapsulasi operasi Milvus, termasuk koneksi, pembuatan koleksi, pembuatan indeks, penyisipan data, dan pencarian.

Prosedur

if __name__ == "__main__":
    # Konfigurasikan API Milvus dan DashScope
    MILVUS_TOKEN = "root:****"
    MILVUS_HOST = "c-0aa16b1****.milvus.aliyuncs.com"
    MILVUS_PORT = "19530"
    COLLECTION_NAME = "multimodal_search"
    INDEX = "IVF_FLAT"  # IVF_FLAT ATAU HNSW  
    script_dir = os.path.dirname(os.path.abspath(__file__))
    csv_path = os.path.join(script_dir, "reverse_image_search.csv")



    # Langkah 1: Inisialisasi klien Milvus
    milvus_client = MilvusClient(MILVUS_TOKEN, MILVUS_HOST, MILVUS_PORT, INDEX, COLLECTION_NAME)

    # Langkah 2: Inisialisasi LLM Qwen-VL dan model penyematan multimodal
    extractor = FeatureExtractor(DASHSCOPE_API_KEY)
    extractorVL = FeatureExtractorVL(DASHSCOPE_API_KEY)

    # Langkah 3: Hasilkan embedding untuk set data gambar dan masukkan ke Milvus
    embeddings = load_image_embeddings(extractor, extractorVL, csv_path)
    milvus_client.insert(embeddings)

Generating image embeddings:  100%
INFO:__main__:Data inserted and loaded successfully.

if __name__ == "__main__":
    MILVUS_HOST = "c-xxxxxxxxxxxx.milvus.aliyuncs.com"
    MILVUS_PORT = "19530"
    MILVUS_TOKEN = "root:****"
    COLLECTION_NAME = "multimodal_search"
    INDEX = "IVF_FLAT" # IVF_FLAT ATAU HNSW
    DASHSCOPE_API_KEY = "<YOUR_DASHSCOPE_API_KEY >"
    
    # Langkah 1: Inisialisasi klien Milvus
    milvus_client = MilvusClient(MILVUS_TOKEN, MILVUS_HOST, MILVUS_PORT, INDEX, COLLECTION_NAME)
    
    # Langkah 2: Inisialisasi model penyematan multimodal
    extractor = FeatureExtractor(DASHSCOPE_API_KEY)

    # Langkah 4: Contoh pencarian multimodal untuk pencarian teks-ke-gambar dan teks-ke-teks
    text_query = "a brown dog"
    text_embedding = extractor(text_query, "text")
    text_results_1 = milvus_client.search(text_embedding, field = 'image_embedding')
    logger.info(f"Text-to-image search results: {text_results_1}")
    text_results_2 = milvus_client.search(text_embedding, field = 'text_embedding')
    logger.info(f"Text-to-text search results: {text_results_2}")
  

INFO:__main__:Text-to-image search results: [
{'id': 456882250782308942, 'distance': 1.338853359222412, 'origin': './train/Rhodesian_ridgeback/n02087394_9675.JPEG', 'image_description': 'A photo of a small dog standing on a carpet. It has brown fur and blue eyes.\nKeywords: puppy, carpet, eyes, fur color, standing'}, 
{'id': 456882250782308933, 'distance': 1.3568601608276367, 'origin': './train/Rhodesian_ridgeback/n02087394_6382.JPEG', 'image_description': 'This is a brown hound with drooping ears and a collar around its neck. It is looking straight ahead.\n\nKeywords: dog, brown, hound, ears, collar'}, 
{'id': 456882250782308940, 'distance': 1.3838427066802979, 'origin': './train/Rhodesian_ridgeback/n02087394_5846.JPEG', 'image_description': 'Two puppies are playing on a blanket. One dog is lying on top of the other, with a teddy bear in the background.\n\nKeywords: puppies, playing, blanket, teddy bear, interaction'}]
INFO:__main__:Text-to-text search results: [
{'id': 456882250782309025, 'distance': 0.6969608068466187, 'origin': './train/mongoose/n02137549_7552.JPEG', 'image_description': 'This is a close-up photo of a small brown animal. It has a round face and large eyes.\n\nKeywords: small animal, brown fur, round face, large eyes, natural background'}, 
{'id': 456882250782308933, 'distance': 0.7110348343849182, 'origin': './train/Rhodesian_ridgeback/n02087394_6382.JPEG', 'image_description': 'This is a brown hound with drooping ears and a collar around its neck. It is looking straight ahead.\n\nKeywords: dog, brown, hound, ears, collar'}, 
{'id': 456882250782308992, 'distance': 0.7725887298583984, 'origin': './train/lion/n02129165_19310.JPEG', 'image_description': 'This is a close-up photo of a lion. It has a thick mane and sharp eyes.\n\nKeywords: lion, eyes, mane, natural environment, wild animal'}]

if __name__ == "__main__":
    # Konfigurasikan API Milvus dan DashScope
    MILVUS_TOKEN = "root:****"
    MILVUS_HOST = "c-0aa16b1****.milvus.aliyuncs.com"
    MILVUS_PORT = "19530"
    COLLECTION_NAME = "multimodal_search"
    INDEX = "IVF_FLAT"  # IVF_FLAT ATAU HNSW
    DASHSCOPE_API_KEY = "<YOUR_DASHSCOPE_API_KEY >"

    # Langkah 1: Inisialisasi klien Milvus
    milvus_client = MilvusClient(MILVUS_TOKEN, MILVUS_HOST, MILVUS_PORT, INDEX, COLLECTION_NAME)
  
    # Langkah 2: Inisialisasi model penyematan multimodal
    extractor = FeatureExtractor(DASHSCOPE_API_KEY)

    # Langkah 5: Contoh pencarian multimodal untuk pencarian berdasarkan gambar dan pencarian gambar-ke-teks
    image_query_path = "./test/lion/n02129165_13728.JPEG"
    image_embedding = extractor(image_query_path, "image")
    image_results_1 = milvus_client.search(image_embedding, field = 'image_embedding')
    logger.info(f"Search by image results: {image_results_1}")
    image_results_2 = milvus_client.search(image_embedding, field = 'text_embedding')
    logger.info(f"Image-to-text search results: {image_results_2}")

INFO:__main__:Search by image results: [
{'id': 456882250782308987, 'distance': 0.23892249166965485, 'origin': './train/lion/n02129165_19953.JPEG', 'image_description': 'A majestic lion stands by a rock, with trees and bushes in the background. Sunlight shines on its body.\n\nKeywords: lion, rock, forest, sunlight, wildness'}, 
{'id': 456882250782308989, 'distance': 0.4113130569458008, 'origin': './train/lion/n02129165_1142.JPEG', 'image_description': 'A lion rests among dense green plants. The background consists of bamboo and trees.\n\nKeywords: lion, grass, green plants, tree trunk, natural environment'}, 
{'id': 456882250782308984, 'distance': 0.5206397175788879, 'origin': './train/lion/n02129165_16.JPEG', 'image_description': 'The image shows a pair of lions standing on the grass. The male lion has a thick mane, while the female lion appears leaner.\n\nKeywords: lion, grass, male, female, natural environment'}]
INFO:__main__:Image-to-text search results: 
[{'id': 456882250782308989, 'distance': 1.0935896635055542, 'origin': './train/lion/n02129165_1142.JPEG', 'image_description': 'A lion rests among dense green plants. The background consists of bamboo and trees.\n\nKeywords: lion, grass, green plants, tree trunk, natural environment'}, 
{'id': 456882250782308987, 'distance': 1.2102885246276855, 'origin': './train/lion/n02129165_19953.JPEG', 'image_description': 'A majestic lion stands by a rock, with trees and bushes in the background. Sunlight shines on its body.\n\nKeywords: lion, rock, forest, sunlight, wildness'}, 
{'id': 456882250782308992, 'distance': 1.2725986242294312, 'origin': './train/lion/n02129165_19310.JPEG', 'image_description': 'This is a close-up photo of a lion. It has a thick mane and sharp eyes.\n\nKeywords: lion, eyes, mane, natural environment, wild animal'}]