Recommandation de films avec Milvus

Dans ce carnet, nous allons explorer comment générer des embeddings de descriptions de films en utilisant OpenAI et exploiter ces embeddings dans Milvus pour recommander des films qui correspondent à vos préférences. Pour améliorer nos résultats de recherche, nous utiliserons le filtrage pour effectuer des recherches de métadonnées. L'ensemble de données utilisé dans cet exemple provient de HuggingFace datasets et contient plus de 8 000 entrées de films, ce qui constitue un riche réservoir d'options pour les recommandations de films.

Dépendances et environnement

Vous pouvez installer les dépendances en exécutant la commande suivante :

$ pip install openai pymilvus datasets tqdm

import os

os.environ["OPENAI_API_KEY"] = "sk-***********"

Initialiser le client OpenAI et Milvus

Initialiser le client OpenAI.

from openai import OpenAI

openai_client = OpenAI()

Définissez le nom et la dimension de la collection pour les embeddings.

COLLECTION_NAME = "movie_search"
DIMENSION = 1536

BATCH_SIZE = 1000

Se connecter à Milvus.

from pymilvus import MilvusClient

# Connect to Milvus Database
client = MilvusClient("./milvus_demo.db")

# Remove collection if it already exists
if client.has_collection(COLLECTION_NAME):
    client.drop_collection(COLLECTION_NAME)

Définissez les champs de la collection, qui comprennent l'identifiant, le titre, le type, l'année de sortie, la note et la description.

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)

Créez l'index sur la collection et chargez-le.

# 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)

Ensemble de données

Milvus étant opérationnel, nous pouvons commencer à saisir nos données. Hugging Face Datasets est un hub qui contient de nombreux ensembles de données d'utilisateurs différents et, pour cet exemple, nous utilisons l'ensemble de données netflix-shows de HuggingLearners. Cet ensemble de données contient des films et leurs paires de métadonnées pour plus de 8 000 films. Nous allons intégrer chaque description et la stocker dans Milvus avec son titre, son type, son année de sortie et sa note.

from datasets import load_dataset

dataset = load_dataset("hugginglearners/netflix-shows", split="train")

Insérer les données

Maintenant que nous avons nos données sur notre machine, nous pouvons commencer à les intégrer et à les insérer dans Milvus. La fonction d'intégration prend du texte et renvoie les intégrations sous forme de liste.

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]

L'étape suivante est celle de l'insertion proprement dite. Nous itérons à travers toutes les entrées et créons des lots que nous insérons une fois que nous avons atteint la taille de lot définie. Une fois la boucle terminée, nous insérons le dernier lot restant s'il existe.

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 = []

Interroger la base de données

Nos données étant insérées en toute sécurité dans Milvus, nous pouvons maintenant effectuer une requête. La requête prend en compte un tuple de la description du film que vous recherchez et le filtre à utiliser. Vous trouverez plus d'informations sur le filtre ici. La recherche affiche d'abord la description et l'expression du filtre. Ensuite, pour chaque résultat, nous affichons le score, le titre, le type, l'année de sortie, la note et la description des films obtenus.

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.