Recomendación de películas con Milvus

En este cuaderno, exploraremos cómo generar incrustaciones de descripciones de películas utilizando OpenAI y aprovechar esas incrustaciones dentro de Milvus para recomendar películas que coincidan con sus preferencias. Para mejorar nuestros resultados de búsqueda, utilizaremos el filtrado para realizar búsquedas de metadatos. El conjunto de datos utilizado en este ejemplo procede de HuggingFace datasets y contiene más de 8.000 entradas de películas, lo que proporciona un amplio abanico de opciones para la recomendación de películas.

Dependencias y entorno

Puedes instalar las dependencias ejecutando el siguiente comando:

$ pip install openai pymilvus datasets tqdm

import os

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

Inicializar el cliente OpenAI y Milvus

Inicialice el cliente OpenAI.

from openai import OpenAI

openai_client = OpenAI()

Establezca el nombre de la colección y la dimensión de las incrustaciones.

COLLECTION_NAME = "movie_search"
DIMENSION = 1536

BATCH_SIZE = 1000

Conéctese a 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)

Defina los campos para la colección, que incluyen el id, el título, el tipo, el año de lanzamiento, la calificación y la descripción.

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)

Crea el índice de la colección y cárgalo.

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

Conjunto de datos

Con Milvus en funcionamiento podemos empezar a obtener nuestros datos. Hugging Face Datasets es un centro que contiene muchos conjuntos de datos de usuario diferentes, y para este ejemplo estamos utilizando el conjunto de datos netflix-shows de HuggingLearners. Este conjunto de datos contiene películas y sus pares de metadatos para más de 8.000 películas. Vamos a incrustar cada descripción y almacenarla en Milvus junto con su título, tipo, año de lanzamiento y calificación.

from datasets import load_dataset

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

Insertar los datos

Ahora que tenemos nuestros datos en nuestra máquina podemos empezar a incrustarlos e insertarlos en Milvus. La función de incrustación toma el texto y devuelve las incrustaciones en un formato de lista.

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]

El siguiente paso es la inserción propiamente dicha. Recorremos todas las entradas y creamos lotes que insertamos una vez que alcanzamos el tamaño de lote establecido. Una vez finalizado el bucle, insertamos el último lote restante si 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 = []

Consultar la base de datos

Con nuestros datos insertados de forma segura en Milvus, ahora podemos realizar una consulta. La consulta toma una tupla de la descripción de la película que está buscando y el filtro a utilizar. Puede encontrar más información sobre el filtro aquí. La búsqueda imprime primero la descripción y la expresión del filtro. A continuación, para cada resultado se imprime la puntuación, el título, el tipo, el año de estreno, la calificación y la descripción de las películas resultantes.

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.