Filmempfehlungen mit Milvus

In diesem Notizbuch werden wir untersuchen, wie wir mit OpenAI Einbettungen von Filmbeschreibungen erzeugen und diese Einbettungen in Milvus nutzen können, um Filme zu empfehlen, die Ihren Vorlieben entsprechen. Um unsere Suchergebnisse zu verbessern, werden wir Filter für die Suche nach Metadaten verwenden. Der Datensatz, der in diesem Beispiel verwendet wird, stammt von HuggingFace Datasets und enthält über 8.000 Filmeinträge, die einen reichen Pool an Optionen für Filmempfehlungen bieten.

Abhängigkeiten und Umgebung

Sie können die Abhängigkeiten installieren, indem Sie den folgenden Befehl ausführen:

$ pip install openai pymilvus datasets tqdm

import os

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

Initialisieren Sie den OpenAI-Client und Milvus

Initialisieren Sie den OpenAI-Client.

from openai import OpenAI

openai_client = OpenAI()

Legen Sie den Sammlungsnamen und die Dimension für die Einbettungen fest.

COLLECTION_NAME = "movie_search"
DIMENSION = 1536

BATCH_SIZE = 1000

Verbinden Sie sich mit 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)

Definieren Sie die Felder für die Sammlung, die die ID, den Titel, den Typ, das Erscheinungsjahr, die Bewertung und die Beschreibung enthalten.

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)

Erstellen Sie den Index für die Sammlung und laden Sie ihn.

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

Datensatz

Wenn Milvus eingerichtet ist und läuft, können wir mit dem Abrufen unserer Daten beginnen. Hugging Face Datasets ist ein Hub, der viele verschiedene Benutzerdatensätze enthält, und für dieses Beispiel verwenden wir den Netflix-Shows-Datensatz von HuggingLearners. Dieser Datensatz enthält Filme und ihre Metadatenpaare für über 8 Tausend Filme. Wir werden jede Beschreibung einbetten und in Milvus zusammen mit dem Titel, dem Typ, dem Erscheinungsjahr und der Bewertung speichern.

from datasets import load_dataset

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

Einfügen der Daten

Nun, da wir unsere Daten auf unserem Rechner haben, können wir damit beginnen, sie einzubetten und in Milvus einzufügen. Die Einbettungsfunktion nimmt Text auf und gibt die Einbettungen in einem Listenformat zurück.

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]

Der nächste Schritt ist das eigentliche Einfügen. Wir durchlaufen alle Einträge und erstellen Stapel, die wir einfügen, sobald wir die festgelegte Stapelgröße erreicht haben. Nach Beendigung der Schleife wird der letzte verbleibende Stapel eingefügt, sofern er existiert.

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

Abfrage der Datenbank

Nachdem wir unsere Daten sicher in Milvus eingefügt haben, können wir nun eine Abfrage durchführen. Die Abfrage nimmt ein Tupel mit der gesuchten Filmbeschreibung und dem zu verwendenden Filter auf. Mehr Informationen über den Filter finden Sie hier. Die Suche gibt zunächst die Beschreibung und den Filterausdruck aus. Danach geben wir für jedes Ergebnis die Punktzahl, den Titel, den Typ, das Erscheinungsjahr, die Bewertung und die Beschreibung der gefundenen Filme aus.

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.