Milvus와 SentenceTransformers를 사용한 영화 검색

이 예에서는 Milvus와 SentenceTransformers 라이브러리를 사용해 Wikipedia 문서 검색을 살펴보겠습니다. 우리가 검색하는 데이터 세트는 Kaggle에 있는 Wikipedia-Movie-Plots 데이터 세트입니다. 이 예제에서는 데이터를 공용 Google 드라이브에 리호스팅했습니다.

시작해 보겠습니다.

설치 요구 사항

이 예제에서는 Milvus를 사용하기 위해 pymilvus, 벡터 임베딩을 생성하기 위해 sentencetransformers, 예제 데이터 집합을 다운로드하기 위해 gdown 에 연결할 것입니다.

pip install pymilvus sentence-transformers gdown

데이터 가져오기

gdown 을 사용하여 Google 드라이브에서 압축 파일을 가져온 다음 기본 제공 zipfile 라이브러리로 압축을 풀겠습니다.

import gdown
url = 'https://drive.google.com/uc?id=11ISS45aO2ubNCGaC3Lvd3D7NT8Y7MeO8'
output = './movies.zip'
gdown.download(url, output)

import zipfile

with zipfile.ZipFile("./movies.zip","r") as zip_ref:
    zip_ref.extractall("./movies")

전역 매개변수

여기에서 자신의 계정으로 실행하기 위해 수정해야 하는 주요 인수를 찾을 수 있습니다. 각 인수 옆에는 해당 인수가 무엇인지에 대한 설명이 있습니다.

# Milvus Setup Arguments
COLLECTION_NAME = 'movies_db'  # Collection name
DIMENSION = 384  # Embeddings size
COUNT = 1000  # Number of vectors to insert
MILVUS_HOST = 'localhost'
MILVUS_PORT = '19530'

# Inference Arguments
BATCH_SIZE = 128

# Search Arguments
TOP_K = 3

Milvus 설정하기

이 시점에서 Milvus 설정을 시작하겠습니다. 단계는 다음과 같습니다:

1. 제공된 URI를 사용하여 Milvus 인스턴스에 연결합니다.

   from pymilvus import connections
   
   # Connect to Milvus Database
   connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
   

2. 컬렉션이 이미 존재하는 경우 삭제합니다.

   from pymilvus import utility
   
   # Remove any previous collections with the same name
   if utility.has_collection(COLLECTION_NAME):
       utility.drop_collection(COLLECTION_NAME)
   

3. 아이디, 동영상 제목, 줄거리 텍스트 임베딩을 포함하는 컬렉션을 생성합니다.

   from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
   
   
   # Create collection which includes the id, title, and embedding.
   fields = [
       FieldSchema(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True),
       FieldSchema(name='title', dtype=DataType.VARCHAR, max_length=200),  # VARCHARS need a maximum length, so for this example they are set to 200 characters
       FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
   ]
   schema = CollectionSchema(fields=fields)
   collection = Collection(name=COLLECTION_NAME, schema=schema)
   

4. 새로 만든 컬렉션에 인덱스를 생성하고 메모리에 로드합니다.

   # Create an IVF_FLAT index for collection.
   index_params = {
       'metric_type':'L2',
       'index_type':"IVF_FLAT",
       'params':{'nlist': 1536}
   }
   collection.create_index(field_name="embedding", index_params=index_params)
   collection.load()
   

이 단계가 완료되면 컬렉션을 삽입하고 검색할 준비가 된 것입니다. 추가된 모든 데이터는 자동으로 색인화되어 즉시 검색할 수 있습니다. 데이터가 매우 새 데이터인 경우, 아직 색인 작업이 진행 중인 데이터에 대해 무차별 대입 검색이 사용되므로 검색 속도가 느려질 수 있습니다.

데이터 삽입하기

이 예제에서는 SentenceTransformers miniLM 모델을 사용하여 플롯 텍스트의 임베딩을 만들겠습니다. 이 모델은 384딤 임베딩을 반환합니다.

다음 몇 단계에서는 다음과 같이 하겠습니다:

1. 데이터 로드하기.
2. SentenceTransformers를 사용하여 플롯 텍스트 데이터 임베딩하기.
3. Milvus에 데이터 삽입하기.

import csv
from sentence_transformers import SentenceTransformer

transformer = SentenceTransformer('all-MiniLM-L6-v2')

# Extract the book titles
def csv_load(file):
    with open(file, newline='') as f:
        reader = csv.reader(f, delimiter=',')
        for row in reader:
            if '' in (row[1], row[7]):
                continue
            yield (row[1], row[7])


# Extract embedding from text using OpenAI
def embed_insert(data):
    embeds = transformer.encode(data[1]) 
    ins = [
            data[0],
            [x for x in embeds]
    ]
    collection.insert(ins)

import time

data_batch = [[],[]]

count = 0

for title, plot in csv_load('./movies/plots.csv'):
    if count <= COUNT:
        data_batch[0].append(title)
        data_batch[1].append(plot)
        if len(data_batch[0]) % BATCH_SIZE == 0:
            embed_insert(data_batch)
            data_batch = [[],[]]
        count += 1
    else:
        break

# Embed and insert the remainder
if len(data_batch[0]) != 0:
    embed_insert(data_batch)

# Call a flush to index any unsealed segments.
collection.flush()

검색 수행하기

Milvus에 모든 데이터를 삽입했으면 검색을 시작할 수 있습니다. 이 예에서는 줄거리를 기반으로 영화를 검색하겠습니다. 일괄 검색을 수행하기 때문에 검색 시간은 영화 검색 전체에 걸쳐 공유됩니다.

# Search for titles that closest match these phrases.
search_terms = ['A movie about cars', 'A movie about monsters']

# Search the database based on input text
def embed_search(data):
    embeds = transformer.encode(data) 
    return [x for x in embeds]

search_data = embed_search(search_terms)

start = time.time()
res = collection.search(
    data=search_data,  # Embeded search value
    anns_field="embedding",  # Search across embeddings
    param={},
    limit = TOP_K,  # Limit to top_k results per search
    output_fields=['title']  # Include title field in result
)
end = time.time()

for hits_i, hits in enumerate(res):
    print('Title:', search_terms[hits_i])
    print('Search Time:', end-start)
    print('Results:')
    for hit in hits:
        print( hit.entity.get('title'), '----', hit.distance)
    print()

결과는 다음과 비슷할 것입니다:

Title: A movie about cars
Search Time: 0.08636689186096191
Results:
Youth's Endearing Charm ---- 1.0954499244689941
From Leadville to Aspen: A Hold-Up in the Rockies ---- 1.1019384860992432
Gentlemen of Nerve ---- 1.1331942081451416

Title: A movie about monsters
Search Time: 0.08636689186096191
Results:
The Suburbanite ---- 1.0666425228118896
Youth's Endearing Charm ---- 1.1072258949279785
The Godless Girl ---- 1.1511223316192627