Rekomendasi Film dengan Milvus

Dalam buku catatan ini, kita akan mengeksplorasi cara membuat embedding deskripsi film menggunakan OpenAI dan memanfaatkan embedding tersebut di dalam Milvus untuk merekomendasikan film yang sesuai dengan preferensi Anda. Untuk meningkatkan hasil pencarian, kita akan menggunakan pemfilteran untuk melakukan pencarian metadata. Dataset yang digunakan dalam contoh ini bersumber dari dataset HuggingFace dan berisi lebih dari 8.000 entri film, memberikan banyak pilihan untuk rekomendasi film.

Ketergantungan dan Lingkungan

Anda dapat menginstal dependensi dengan menjalankan perintah berikut:

$ pip install openai pymilvus datasets tqdm

import os

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

Inisialisasi klien OpenAI dan Milvus

Inisialisasi klien OpenAI.

from openai import OpenAI

openai_client = OpenAI()

Tetapkan nama koleksi dan dimensi untuk penyematan.

COLLECTION_NAME = "movie_search"
DIMENSION = 1536

BATCH_SIZE = 1000

Hubungkan ke 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)

Tentukan bidang untuk koleksi, yang meliputi id, judul, jenis, tahun rilis, peringkat, dan deskripsi.

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)

Buat indeks pada koleksi dan muat.

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

Dataset

Dengan Milvus aktif dan berjalan, kita dapat mulai mengambil data kita. Hugging Face Datasets adalah hub yang menampung banyak dataset pengguna yang berbeda, dan untuk contoh ini kita menggunakan dataset netflix-shows milik HuggingLearners. Dataset ini berisi film dan pasangan metadatanya untuk lebih dari 8 ribu film. Kita akan menyematkan setiap deskripsi dan menyimpannya di dalam Milvus bersama dengan judul, jenis, tahun rilis, dan peringkatnya.

from datasets import load_dataset

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

Memasukkan Data

Sekarang kita sudah memiliki data di mesin kita, kita bisa mulai menyematkannya dan memasukkannya ke dalam Milvus. Fungsi penyematan mengambil teks dan mengembalikan penyematan dalam format daftar.

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]

Langkah selanjutnya adalah melakukan penyisipan yang sebenarnya. Kita mengulang semua entri dan membuat batch yang kita sisipkan setelah kita mencapai ukuran batch yang kita tetapkan. Setelah perulangan selesai, kita menyisipkan batch remaning terakhir jika ada.

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

Menanyakan Basis Data

Dengan data yang telah dimasukkan ke dalam Milvus, kita sekarang dapat melakukan query. Kueri ini mengambil tuple dari deskripsi film yang Anda cari dan filter yang akan digunakan. Info lebih lanjut tentang filter dapat ditemukan di sini. Pencarian pertama-tama mencetak deskripsi dan ekspresi filter Anda. Setelah itu untuk setiap hasil, kami mencetak skor, judul, jenis, tahun rilis, peringkat dan deskripsi film hasil pencarian.

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.