使用 Zilliz 和 OpenAI 进行筛选搜索

寻找您的下一部电影

在这个笔记本中，我们将讨论如何使用 OpenAI 生成电影描述的嵌入向量，并在 Zilliz 中使用这些嵌入向量来查找相关的电影。为了缩小搜索结果的范围并尝试一些新东西，我们将使用过滤器进行元数据搜索。本示例中的数据集来自 HuggingFace 数据集，包含8000多个电影条目。

让我们首先下载本笔记本所需的库： - openai 用于与 OpenAI 嵌入服务进行通信 - pymilvus 用于与 Zilliz 服务器进行通信 - datasets 用于下载数据集 - tqdm 用于显示进度条

! pip install openai pymilvus datasets tqdm

要使Zilliz运行起来，请查看这里。在设置好您的账户和数据库之后，继续设置以下数值： - URI：您的数据库运行的URI - USER：您的数据库用户名 - PASSWORD：您的数据库密码 - COLLECTION_NAME：在Zilliz中命名集合的名称 - DIMENSION：嵌入的维度 - OPENAI_ENGINE：要使用的嵌入模型 - openai.api_key：您的OpenAI账户密钥 - INDEX_PARAM：用于集合的索引设置 - QUERY_PARAM：要使用的搜索参数 - BATCH_SIZE：一次要嵌入和插入多少个文本

import openai

URI = 'your_uri'
TOKEN = 'your_token' # TOKEN == 用户:密码 或 api_key
COLLECTION_NAME = 'book_search'
DIMENSION = 1536
OPENAI_ENGINE = 'text-embedding-3-small'
openai.api_key = 'sk-your_key'

INDEX_PARAM = {
    'metric_type':'L2',
    'index_type':"AUTOINDEX",
    'params':{}
}

QUERY_PARAM = {
    "metric_type": "L2",
    "params": {},
}

BATCH_SIZE = 1000

from pymilvus import connections, utility, FieldSchema, Collection, CollectionSchema, DataType

# 连接到Zilliz数据库
connections.connect(uri=URI, token=TOKEN)

# 如果集合已存在，则将其移除。
if utility.has_collection(COLLECTION_NAME):
    utility.drop_collection(COLLECTION_NAME)

# 创建一个集合，包含id、标题和嵌入信息。
fields = [
    FieldSchema(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name='title', dtype=DataType.VARCHAR, max_length=64000),
    FieldSchema(name='type', dtype=DataType.VARCHAR, max_length=64000),
    FieldSchema(name='release_year', dtype=DataType.INT64),
    FieldSchema(name='rating', dtype=DataType.VARCHAR, max_length=64000),
    FieldSchema(name='description', dtype=DataType.VARCHAR, max_length=64000),
    FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]
schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)

# 在集合上创建索引并加载它。
collection.create_index(field_name="embedding", index_params=INDEX_PARAM)
collection.load()

数据集

有了Zilliz运行起来，我们就可以开始获取我们的数据了。Hugging Face Datasets 是一个包含许多不同用户数据集的中心，而在这个示例中，我们使用了HuggingLearners的netflix-shows数据集。该数据集包含超过8000部电影的电影及其元数据对。我们将嵌入每个描述并将其与标题、类型、发布年份和评分一起存储在Zilliz中。

import datasets

# 下载数据集 
dataset = datasets.load_dataset('hugginglearners/netflix-shows', split='train')

/Users/filiphaltmayer/miniconda3/envs/haystack/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
Found cached dataset csv (/Users/filiphaltmayer/.cache/huggingface/datasets/hugginglearners___csv/hugginglearners--netflix-shows-03475319fc65a05a/0.0.0/6b34fb8fcf56f7c8ba51dc895bfa2bfbe43546f190a60fcf74bb5e8afdcc2317)

插入数据

现在我们已经将数据保存在我们的机器上，我们可以开始将其嵌入并插入到 Zilliz 中。嵌入函数接受文本并以列表格式返回嵌入。

# 简单函数，用于将文本转换为嵌入表示。
def embed(texts):
    embeddings = openai.Embedding.create(
        input=texts,
        engine=OPENAI_ENGINE
    )
    return [x['embedding'] for x in embeddings['data']]

接下来的步骤是实际的插入操作。我们遍历所有条目，并创建批次，一旦达到设定的批次大小，我们就插入这些批次。循环结束后，如果存在剩余的最后一个批次，我们也会将其插入。

from tqdm import tqdm

data = [
    [], # 标题
    [], # 类型
    [], # 发行年份
    [], # 评级
    [], # 描述
]

# 批量嵌入和插入
for i in tqdm(range(0, len(dataset))):
    data[0].append(dataset[i]['title'] or '')
    data[1].append(dataset[i]['type'] or '')
    data[2].append(dataset[i]['release_year'] or -1)
    data[3].append(dataset[i]['rating'] or '')
    data[4].append(dataset[i]['description'] or '')
    if len(data[0]) % BATCH_SIZE == 0:
        data.append(embed(data[4]))
        collection.insert(data)
        data = [[],[],[],[],[]]

# 嵌入并插入余数 
if len(data[0]) != 0:
    data.append(embed(data[4]))
    collection.insert(data)
    data = [[],[],[],[],[]]

100%|██████████| 8807/8807 [00:54<00:00, 162.59it/s]

查询数据库

在我们的数据安全地插入到 Zilliz 后，现在我们可以执行查询操作了。查询操作接受一个包含您要搜索的电影描述和要使用的过滤器的元组。有关过滤器的更多信息可以在这里找到。搜索首先打印出您的描述和过滤器表达式。然后对于每个结果，我们打印出电影的得分、标题、类型、发行年份、评分和描述。

import textwrap

def query(query, top_k = 5):
    text, expr = query
    res = collection.search(embed(text), anns_field='embedding', expr = expr, param=QUERY_PARAM, limit = top_k, output_fields=['title', 'type', 'release_year', 'rating', 'description'])
    for i, hit in enumerate(res):
        print('Description:', text, 'Expression:', expr)
        print('Results:')
        for ii, hits in enumerate(hit):
            print('\t' + 'Rank:', ii + 1, 'Score:', hits.score, 'Title:', hits.entity.get('title'))
            print('\t\t' + 'Type:', hits.entity.get('type'), 'Release Year:', hits.entity.get('release_year'), 'Rating:', hits.entity.get('rating'))
            print(textwrap.fill(hits.entity.get('description'), 88))
            print()

my_query = ('movie about a fluffly animal', 'release_year < 2019 and rating like \"PG%\"')

query(my_query)

Description: movie about a fluffly animal Expression: release_year < 2019 and rating like "PG%"
Results:
    Rank: 1 Score: 0.30085673928260803 Title: The Lamb
        Type: Movie Release Year: 2017 Rating: PG
A big-dreaming donkey escapes his menial existence and befriends some free-spirited
animal pals in this imaginative retelling of the Nativity Story.

    Rank: 2 Score: 0.3352621793746948 Title: Puss in Boots
        Type: Movie Release Year: 2011 Rating: PG
The fabled feline heads to the Land of Giants with friends Humpty Dumpty and Kitty
Softpaws on a quest to nab its greatest treasure: the Golden Goose.

    Rank: 3 Score: 0.3415083587169647 Title: Show Dogs
        Type: Movie Release Year: 2018 Rating: PG
A rough and tough police dog must go undercover with an FBI agent as a prim and proper
pet at a dog show to save a baby panda from an illegal sale.

    Rank: 4 Score: 0.3428957462310791 Title: Open Season 2
        Type: Movie Release Year: 2008 Rating: PG
Elliot the buck and his forest-dwelling cohorts must rescue their dachshund pal from
some spoiled pets bent on returning him to domesticity.

    Rank: 5 Score: 0.34376364946365356 Title: Stuart Little 2
        Type: Movie Release Year: 2002 Rating: PG
Zany misadventures are in store as lovable city mouse Stuart and his human brother,
George, raise the roof in this sequel to the 1999 blockbuster.