使用 Milvus 和 SentenceTransformers 进行电影搜索

在本示例中，我们将使用 Milvus 和 SentenceTransformers 库进行维基百科文章搜索。我们搜索的数据集是Kaggle 上的维基百科-电影-情节数据集。在本示例中，我们将数据重新托管在谷歌公共驱动器中。

让我们开始吧。

安装要求

在本例中，我们将使用pymilvus 连接使用 Milvus，使用sentencetransformers 生成向量嵌入，使用gdown 下载示例数据集。

pip install pymilvus sentence-transformers gdown

抓取数据

我们将使用gdown 从 Google Drive 抓取压缩包，然后使用内置的zipfile 库解压。

import gdown
url = 'https://drive.google.com/uc?id=11ISS45aO2ubNCGaC3Lvd3D7NT8Y7MeO8'
output = './movies.zip'
gdown.download(url, output)

import zipfile

with zipfile.ZipFile("./movies.zip","r") as zip_ref:
    zip_ref.extractall("./movies")

全局参数

在这里，我们可以找到使用自己的账户运行时需要修改的主要参数。每个参数旁边都有说明。

# Milvus Setup Arguments
COLLECTION_NAME = 'movies_db'  # Collection name
DIMENSION = 384  # Embeddings size
COUNT = 1000  # Number of vectors to insert
MILVUS_HOST = 'localhost'
MILVUS_PORT = '19530'

# Inference Arguments
BATCH_SIZE = 128

# Search Arguments
TOP_K = 3

设置 Milvus

现在，我们开始设置 Milvus。步骤如下：

1. 使用提供的 URI 连接到 Milvus 实例。

   from pymilvus import connections
   
   # Connect to Milvus Database
   connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
   

2. 如果集合已经存在，则删除它。

   from pymilvus import utility
   
   # Remove any previous collections with the same name
   if utility.has_collection(COLLECTION_NAME):
       utility.drop_collection(COLLECTION_NAME)
   

3. 创建包含 id、电影标题和情节文本嵌入的集合。

   from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
   
   
   # Create collection which includes the id, title, and embedding.
   fields = [
       FieldSchema(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True),
       FieldSchema(name='title', dtype=DataType.VARCHAR, max_length=200),  # VARCHARS need a maximum length, so for this example they are set to 200 characters
       FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
   ]
   schema = CollectionSchema(fields=fields)
   collection = Collection(name=COLLECTION_NAME, schema=schema)
   

4. 在新创建的集合上创建索引，并将其加载到内存中。

   # Create an IVF_FLAT index for collection.
   index_params = {
       'metric_type':'L2',
       'index_type':"IVF_FLAT",
       'params':{'nlist': 1536}
   }
   collection.create_index(field_name="embedding", index_params=index_params)
   collection.load()
   

完成这些步骤后，就可以插入并搜索该集合了。任何添加的数据都将自动编制索引，并立即可供搜索。如果数据非常新，搜索可能会慢一些，因为将对仍在编制索引过程中的数据使用暴力搜索。

插入数据

在本例中，我们将使用 SentenceTransformers miniLM 模型来创建情节文本的嵌入。该模型可返回 384 维嵌入。

在接下来的几个步骤中，我们将

1. 加载数据。
2. 使用 SentenceTransformers 嵌入情节文本数据。
3. 将数据插入 Milvus。

import csv
from sentence_transformers import SentenceTransformer

transformer = SentenceTransformer('all-MiniLM-L6-v2')

# Extract the book titles
def csv_load(file):
    with open(file, newline='') as f:
        reader = csv.reader(f, delimiter=',')
        for row in reader:
            if '' in (row[1], row[7]):
                continue
            yield (row[1], row[7])


# Extract embedding from text using OpenAI
def embed_insert(data):
    embeds = transformer.encode(data[1]) 
    ins = [
            data[0],
            [x for x in embeds]
    ]
    collection.insert(ins)

import time

data_batch = [[],[]]

count = 0

for title, plot in csv_load('./movies/plots.csv'):
    if count <= COUNT:
        data_batch[0].append(title)
        data_batch[1].append(plot)
        if len(data_batch[0]) % BATCH_SIZE == 0:
            embed_insert(data_batch)
            data_batch = [[],[]]
        count += 1
    else:
        break

# Embed and insert the remainder
if len(data_batch[0]) != 0:
    embed_insert(data_batch)

# Call a flush to index any unsealed segments.
collection.flush()

执行搜索

将所有数据插入 Milvus 后，我们就可以开始执行搜索了。在本例中，我们将根据情节搜索电影。由于我们进行的是批量搜索，因此搜索时间将在电影搜索中共享。

# Search for titles that closest match these phrases.
search_terms = ['A movie about cars', 'A movie about monsters']

# Search the database based on input text
def embed_search(data):
    embeds = transformer.encode(data) 
    return [x for x in embeds]

search_data = embed_search(search_terms)

start = time.time()
res = collection.search(
    data=search_data,  # Embeded search value
    anns_field="embedding",  # Search across embeddings
    param={},
    limit = TOP_K,  # Limit to top_k results per search
    output_fields=['title']  # Include title field in result
)
end = time.time()

for hits_i, hits in enumerate(res):
    print('Title:', search_terms[hits_i])
    print('Search Time:', end-start)
    print('Results:')
    for hit in hits:
        print( hit.entity.get('title'), '----', hit.distance)
    print()

输出结果应该与下面类似：

Title: A movie about cars
Search Time: 0.08636689186096191
Results:
Youth's Endearing Charm ---- 1.0954499244689941
From Leadville to Aspen: A Hold-Up in the Rockies ---- 1.1019384860992432
Gentlemen of Nerve ---- 1.1331942081451416

Title: A movie about monsters
Search Time: 0.08636689186096191
Results:
The Suburbanite ---- 1.0666425228118896
Youth's Endearing Charm ---- 1.1072258949279785
The Godless Girl ---- 1.1511223316192627