与 SentenceTransformers 集成（integrate_with_sentencetransformers）

Milvus和SentenceTransformers库

在这个示例中，我们将使用Milvus和SentenceTransformers库来进行维基百科文章的搜索。

我们要搜索的数据集是在Kaggle (opens in a new tab)上找到的维基百科电影情节数据集。

在此示例中，我们在公共google drive上重新托管了数据。

让我们开始吧。

安装依赖包

对于这个示例，我们将使用pymilvus 来连接和使用Milvus，sentencetransformers来生成向量嵌入，gdown 用于下载示例数据集。

pip install pymilvus sentence-transformers gdown

获取数据

我们将使用gdown从Google Drive获取zip文件，然后使用内置的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
 
# 连接到 Milvus 数据库
connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
  1. 如果集合已经存在，则删除它。
from pymilvus import utility
 
# 删除具有相同名称的以前集合
if utility.has_collection(COLLECTION_NAME):
    utility.drop_collection(COLLECTION_NAME)
  1. 创建包含 ID、电影标题和情节文本嵌入的集合。
from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
 
# 创建包含 ID、标题和情节文本嵌入的集合
fields = [
    FieldSchema(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name='title', dtype=DataType.VARCHAR, max_length=200),  # VARCHAR 需要一个最大长度，所以为了这个例子，它们被设置为200个字符。
    FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]
schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)
  1. 为新创建的集合创建索引并将其加载到内存中。
#为集合创建IVF_FLAT索引
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维度的嵌入。

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

  • 加载数据。
  • 使用SentenceTransformers来嵌入情节文本数据。
  • 将数据插入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 embeding 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()

上述操作由于嵌入操作需要时间，所以比较耗时。为了将耗时控制在可接受的水平，可以尝试将 全局参数 中的 COUNT 设置为一个合适的值。休息一下，喝杯咖啡吧！

执行搜索

所有数据都插入到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
CLI 概述 (cli_overview)