在本笔记本中，我们将介绍如何使用 OpenAI 生成电影描述的嵌入向量，并在 Zilliz 中使用这些嵌入向量来查找相关电影。为了缩小我们的搜索结果并尝试一些新的东西，我们将使用过滤来进行元数据搜索。此示例中的数据集来源于 HuggingFace datasets，包含 8 千多部电影条目。
首先，让我们开始下载此笔记本所需的库
openai 用于与 OpenAI 嵌入服务通信pymilvus 用于与 Zilliz 服务器通信datasets 用于下载数据集tqdm 用于进度条! pip install openai pymilvus datasets tqdm要启动并运行 Zilliz，请查看此处。设置好您的帐户和数据库后，继续设置以下值
import openai
URI = 'your_uri'
TOKEN = 'your_token' # TOKEN == user:password or 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 = 1000from pymilvus import connections, utility, FieldSchema, Collection, CollectionSchema, DataType
# Connect to Zilliz Database
connections.connect(uri=URI, token=TOKEN)# Remove collection if it already exists
if utility.has_collection(COLLECTION_NAME):
utility.drop_collection(COLLECTION_NAME)# 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=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)# Create the index on the collection and load it.
collection.create_index(field_name="embedding", index_params=INDEX_PARAM)
collection.load()Zilliz 启动并运行后，我们可以开始抓取数据。Hugging Face Datasets 是一个托管许多不同用户数据集的中心，在此示例中，我们使用 HuggingLearners 的 netflix-shows 数据集。此数据集包含 8 千多部电影的电影及其元数据对。我们将嵌入每个描述并将其与其标题、类型、release_year 和评分一起存储在 Zilliz 中。
import datasets
# Download the dataset
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 中。嵌入函数接受文本输入，并以列表格式返回嵌入向量。
# Simple function that converts the texts to embeddings
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 = [
[], # title
[], # type
[], # release_year
[], # rating
[], # description
]
# Embed and insert in batches
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 = [[],[],[],[],[]]
# Embed and insert the remainder
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.