Vector Retrieval Service:DashVector + ModelScope: Implement multi-modal searches

from modelscope.msdatasets import MsDataset

dataset = MsDataset.load("muge", split="validation")

import torch
from modelscope.utils.constant import Tasks
from modelscope.pipelines import pipeline
from modelscope.msdatasets import MsDataset
from dashvector import Client, Doc, DashVectorException, DashVectorCode
from PIL import Image
import base64
import io


def image2str(image):
    image_byte_arr = io.BytesIO()
    image.save(image_byte_arr, format='PNG')
    image_bytes = image_byte_arr.getvalue()
    return base64.b64encode(image_bytes).decode()


if __name__ == '__main__':
    # Initialize the DashVector client.
    client = Client(
      api_key='{your-dashvector-api-key}',
      endpoint='{your-dashvector-cluster-endpoint}'
    )

    # Create a collection by specifying the collection name and the number of vector dimensions. In the Huge model of ChineseCLIP, the number of vector dimensions is 1,024.
    rsp = client.create('muge_embedding', 1024)
    if not rsp:
        raise DashVectorException(rsp.code, reason=rsp.message)

    # Batch generate and store image embeddings.
    collection = client.get('muge_embedding')
    pipe = pipeline(task=Tasks.multi_modal_embedding,
                    model='damo/multi-modal_clip-vit-huge-patch14_zh', 
                    model_revision='v1.0.0')
    ds = MsDataset.load("muge", split="validation")

    BATCH_COUNT = 10
    TOTAL_DATA_NUM = len(ds)
    print(f"Start indexing muge validation data, total data size: {TOTAL_DATA_NUM}, batch size:{BATCH_COUNT}")
    idx = 0
    while idx < TOTAL_DATA_NUM:
        batch_range = range(idx, idx + BATCH_COUNT) if idx + BATCH_COUNT <= TOTAL_DATA_NUM else range(idx, TOTAL_DATA_NUM)
        images = [ds[i]['image'] for i in batch_range]
        # Generate image embeddings by using the ChineseCLIP model.
        image_embeddings = pipe.forward({'img': images})['img_embedding']
        image_vectors = image_embeddings.detach().cpu().numpy()
        collection.insert(
            [
                Doc(
                    id=str(img_id),
                    vector=img_vec,
                    fields={'png_img': image2str(img)}
                )
                for img_id, img_vec, img in zip(batch_range, image_vectors, images)
            ]
        )
        idx += BATCH_COUNT
    print("Finish indexing muge validation data")

import torch
from modelscope.utils.constant import Tasks
from modelscope.pipelines import pipeline
from modelscope.msdatasets import MsDataset
from dashvector import Client, Doc, DashVectorException
from PIL import Image
import base64
import io


def str2image(image_str):
    image_bytes = base64.b64decode(image_str)
    return Image.open(io.BytesIO(image_bytes))


def multi_modal_search(input_text):
    # Initialize the DashVector client.
    client = Client(
      api_key='{your-dashvector-api-key}',
      endpoint='{your-dashvector-cluster-endpoint}'
    )

    # Obtain the collection that stores the relevant embeddings.
    collection = client.get('muge_embedding')

    # Obtain the embedding of the query text.
    pipe = pipeline(task=Tasks.multi_modal_embedding,
                    model='damo/multi-modal_clip-vit-huge-patch14_zh', model_revision='v1.0.0')
    text_embedding = pipe.forward({'text': input_text})['text_embedding']  # 2D Tensor, [The number of words in the text, Feature dimensions]
    text_vector = text_embedding.detach().cpu().numpy()[0]

    # Search in DashVector.
    rsp = collection.query(text_vector, topk=3)
    image_list = list()
    for doc in rsp:
        image_str = doc.fields['png_img']
        image_list.append(str2image(image_str))
    return image_list


if __name__ == '__main__':
    text_query = "戴眼镜的狗"
    
    images = multi_modal_search(text_query)
    for img in images:
        # Note: You may need to install an image viewer on the Linux server for the show() function to work.
        # We recommend that you run the code on a server that supports Jupyter Notebook.
        img.show()