Ricerca di film con Milvus e SentenceTransformers

In questo esempio, esamineremo una ricerca di articoli di Wikipedia utilizzando Milvus e la libreria SentenceTransformers. Il set di dati che stiamo cercando è il Wikipedia-Movie-Plots Dataset trovato su Kaggle. Per questo esempio, abbiamo rehosted i dati in un google drive pubblico.

Iniziamo.

Requisiti per l'installazione

Per questo esempio, utilizzeremo pymilvus per connetterci a Milvus, sentencetransformers per generare embeddings vettoriali e gdown per scaricare il dataset di esempio.

pip install pymilvus sentence-transformers gdown

Prelevare i dati

Utilizzeremo gdown per prelevare lo zip da Google Drive e poi decomprimerlo con la libreria integrata 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")

Parametri globali

Qui si trovano gli argomenti principali che devono essere modificati per l'esecuzione con i propri account. Accanto a ciascuno di essi è riportata una descrizione.

# 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

Impostazione di Milvus

A questo punto, iniziamo a configurare Milvus. I passaggi sono i seguenti:

1. Collegarsi all'istanza di Milvus utilizzando l'URI fornito.

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

2. Se la collezione esiste già, eliminarla.

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

3. Creare la collezione che contiene l'id, il titolo del film e le incorporazioni del testo della trama.

   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. Creare un indice sulla collezione appena creata e caricarla in memoria.

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

Una volta eseguiti questi passaggi, la raccolta è pronta per essere inserita e ricercata. Tutti i dati aggiunti verranno indicizzati automaticamente e saranno immediatamente disponibili per la ricerca. Se i dati sono molto recenti, la ricerca potrebbe essere più lenta, in quanto la ricerca brute force verrà utilizzata sui dati ancora in fase di indicizzazione.

Inserimento dei dati

Per questo esempio, utilizzeremo il modello SentenceTransformers miniLM per creare embeddings del testo della trama. Questo modello restituisce embeddings a 384 dimensioni.

Nei prossimi passi si procederà a:

1. Caricare i dati.
2. Incorporare i dati del testo della trama usando SentenceTransformers.
3. Inserire i dati in 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()

Esecuzione della ricerca

Con tutti i dati inseriti in Milvus, possiamo iniziare a eseguire le nostre ricerche. In questo esempio, cercheremo i film in base alla trama. Poiché stiamo eseguendo una ricerca in batch, il tempo di ricerca è condiviso tra le ricerche dei film.

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

L'output dovrebbe essere simile al seguente:

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