Рекомендация фильмов с помощью Milvus
В этом блокноте мы рассмотрим, как генерировать вкрапления описаний фильмов с помощью OpenAI и использовать эти вкрапления в Milvus для рекомендации фильмов, соответствующих вашим предпочтениям. Чтобы улучшить результаты поиска, мы будем использовать фильтрацию для поиска по метаданным. Набор данных, используемый в этом примере, взят из HuggingFace datasets и содержит более 8 000 записей о фильмах, что обеспечивает богатый выбор вариантов для рекомендаций фильмов.
Зависимости и окружение
Вы можете установить зависимости, выполнив следующую команду:
$ pip install openai pymilvus datasets tqdm
Если вы используете Google Colab, то для включения только что установленных зависимостей вам может потребоваться перезапустить среду выполнения (нажмите на меню "Runtime" в верхней части экрана и выберите "Restart session" из выпадающего меню).
В этом примере мы будем использовать OpenAI в качестве LLM. Вам следует подготовить api ключ OPENAI_API_KEY
в качестве переменной окружения.
import os
os.environ["OPENAI_API_KEY"] = "sk-***********"
Инициализация клиента OpenAI и Milvus
Инициализируйте клиент OpenAI.
from openai import OpenAI
openai_client = OpenAI()
Задайте имя коллекции и размер для вкраплений.
COLLECTION_NAME = "movie_search"
DIMENSION = 1536
BATCH_SIZE = 1000
Подключитесь к Milvus.
from pymilvus import MilvusClient
# Connect to Milvus Database
client = MilvusClient("./milvus_demo.db")
Как и в случае с аргументами url
и token
:
- Установка
uri
в качестве локального файла, например./milvus.db
, является наиболее удобным методом, так как он автоматически использует Milvus Lite для хранения всех данных в этом файле. - Если у вас большой объем данных, скажем, более миллиона векторов, вы можете настроить более производительный сервер Milvus на Docker или Kubernetes. В этом случае используйте адрес и порт сервера в качестве uri, например,
http://localhost:19530
. Если вы включили функцию аутентификации на Milvus, используйте "<ваше_имя_пользователя>:<ваш_пароль>" в качестве токена, в противном случае не задавайте токен. - Если вы хотите использовать Zilliz Cloud, полностью управляемый облачный сервис для Milvus, настройте
uri
иtoken
, которые соответствуют публичной конечной точке и ключу Api в Zilliz Cloud.
# Remove collection if it already exists
if client.has_collection(COLLECTION_NAME):
client.drop_collection(COLLECTION_NAME)
Определите поля для коллекции, которые включают в себя id, название, тип, год выпуска, рейтинг и описание.
from pymilvus import DataType
# Create collection which includes the id, title, and embedding.
# 1. Create schema
schema = MilvusClient.create_schema(
auto_id=True,
enable_dynamic_field=False,
)
# 2. Add fields to schema
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="title", datatype=DataType.VARCHAR, max_length=64000)
schema.add_field(field_name="type", datatype=DataType.VARCHAR, max_length=64000)
schema.add_field(field_name="release_year", datatype=DataType.INT64)
schema.add_field(field_name="rating", datatype=DataType.VARCHAR, max_length=64000)
schema.add_field(field_name="description", datatype=DataType.VARCHAR, max_length=64000)
schema.add_field(field_name="embedding", datatype=DataType.FLOAT_VECTOR, dim=DIMENSION)
# 3. Create collection with the schema
client.create_collection(collection_name=COLLECTION_NAME, schema=schema)
Создайте индекс для коллекции и загрузите его.
# Create the index on the collection and load it.
# 1. Prepare index parameters
index_params = client.prepare_index_params()
# 2. Add an index on the embedding field
index_params.add_index(
field_name="embedding", metric_type="IP", index_type="AUTOINDEX", params={}
)
# 3. Create index
client.create_index(collection_name=COLLECTION_NAME, index_params=index_params)
# 4. Load collection
client.load_collection(collection_name=COLLECTION_NAME, replica_number=1)
Набор данных
Запустив Milvus, мы можем приступить к захвату данных. Hugging Face Datasets
- это хаб, в котором хранится множество различных пользовательских наборов данных, и для этого примера мы используем набор данных netflix-shows от HuggingLearners. В этом наборе содержатся фильмы и пары метаданных к ним для более чем 8 тысяч фильмов. Мы собираемся вставить каждое описание и хранить его в Milvus вместе с названием, типом, годом выпуска и рейтингом.
from datasets import load_dataset
dataset = load_dataset("hugginglearners/netflix-shows", split="train")
Вставка данных
Теперь, когда у нас есть данные на машине, мы можем приступить к их встраиванию и вставке в Milvus. Функция встраивания принимает текст и возвращает вложенные данные в виде списка.
def emb_texts(texts):
res = openai_client.embeddings.create(input=texts, model="text-embedding-3-small")
return [res_data.embedding for res_data in res.data]
На следующем этапе происходит собственно вставка. Мы перебираем все записи и создаем партии, которые вставляем, как только достигнем заданного размера партии. После завершения цикла мы вставляем последнюю оставшуюся партию, если она существует.
from tqdm import tqdm
# batch (data to be inserted) is a list of dictionaries
batch = []
# Embed and insert in batches
for i in tqdm(range(0, len(dataset))):
batch.append(
{
"title": dataset[i]["title"] or "",
"type": dataset[i]["type"] or "",
"release_year": dataset[i]["release_year"] or -1,
"rating": dataset[i]["rating"] or "",
"description": dataset[i]["description"] or "",
}
)
if len(batch) % BATCH_SIZE == 0 or i == len(dataset) - 1:
embeddings = emb_texts([item["description"] for item in batch])
for item, emb in zip(batch, embeddings):
item["embedding"] = emb
client.insert(collection_name=COLLECTION_NAME, data=batch)
batch = []
Запрос к базе данных
Теперь, когда наши данные надежно вставлены в Milvus, мы можем выполнить запрос. Запрос содержит кортеж описания фильма, который вы ищете, и фильтр, который нужно использовать. Более подробную информацию о фильтре можно найти здесь. Сначала поиск выводит описание и выражение фильтра. После этого для каждого результата мы выводим оценку, название, тип, год выпуска, рейтинг и описание фильма.
import textwrap
def query(query, top_k=5):
text, expr = query
res = client.search(
collection_name=COLLECTION_NAME,
data=emb_texts(text),
filter=expr,
limit=top_k,
output_fields=["title", "type", "release_year", "rating", "description"],
search_params={
"metric_type": "IP",
"params": {},
},
)
print("Description:", text, "Expression:", expr)
for hit_group in res:
print("Results:")
for rank, hit in enumerate(hit_group, start=1):
entity = hit["entity"]
print(
f"\tRank: {rank} Score: {hit['distance']:} Title: {entity.get('title', '')}"
)
print(
f"\t\tType: {entity.get('type', '')} "
f"Release Year: {entity.get('release_year', '')} "
f"Rating: {entity.get('rating', '')}"
)
description = entity.get("description", "")
print(textwrap.fill(description, width=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.42213767766952515 Title: The Adventures of Tintin
Type: Movie Release Year: 2011 Rating: PG
This 3-D motion capture adapts Georges Remi's classic comic strip about the adventures
of fearless young journalist Tintin and his trusty dog, Snowy.
Rank: 2 Score: 0.4041026830673218 Title: Hedgehogs
Type: Movie Release Year: 2016 Rating: PG
When a hedgehog suffering from memory loss forgets his identity, he ends up on a big
city journey with a pigeon to save his habitat from a human threat.
Rank: 3 Score: 0.3980264663696289 Title: Osmosis Jones
Type: Movie Release Year: 2001 Rating: PG
Peter and Bobby Farrelly outdo themselves with this partially animated tale about an
out-of-shape 40-year-old man who's the host to various organisms.
Rank: 4 Score: 0.39479154348373413 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: 5 Score: 0.39370301365852356 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.