Milvus로 영화 추천하기

이 노트북에서는 OpenAI를 사용해 영화 설명의 임베딩을 생성하고 Milvus 내에서 이러한 임베딩을 활용하여 사용자의 취향에 맞는 영화를 추천하는 방법을 살펴봅니다. 검색 결과를 향상시키기 위해 필터링을 활용하여 메타데이터 검색을 수행합니다. 이 예제에 사용된 데이터 세트는 HuggingFace 데이터 세트에서 제공되며 8,000개 이상의 영화 항목을 포함하고 있어 영화 추천을 위한 풍부한 옵션 풀을 제공합니다.

종속성 및 환경

다음 명령을 실행하여 종속 요소를 설치할 수 있습니다:

$ pip install openai pymilvus datasets tqdm

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

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

아이디, 제목, 유형, 출시 연도, 등급, 설명 등 컬렉션의 필드를 정의합니다.

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 은 다양한 사용자 데이터 세트를 보유하고 있는 허브로, 이 예제에서는 HuggingLearners의 netflix-shows 데이터 세트를 사용합니다. 이 데이터 세트에는 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.