Ricerca di film con Milvus e SentenceTransformers

In questo esempio, cercheremo i riassunti delle trame dei film usando Milvus e la libreria SentenceTransformers. Il dataset che utilizzeremo è Wikipedia Movie Plots with Summaries, ospitato su HuggingFace.

Iniziamo!

Librerie necessarie

Per questo esempio, useremo pymilvus per connetterci a Milvus, sentence-transformers per generare embeddings vettoriali e datasets per scaricare il dataset di esempio.

pip install pymilvus sentence-transformers datasets tqdm

from datasets import load_dataset
from pymilvus import MilvusClient
from pymilvus import FieldSchema, CollectionSchema, DataType
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

Definiremo alcuni parametri globali,

embedding_dim = 384
collection_name = "movie_embeddings"

Scaricare e aprire il set di dati

Con una sola riga, datasets ci permette di scaricare e aprire un set di dati. La libreria metterà in cache il dataset a livello locale e utilizzerà quella copia la prossima volta che verrà eseguita. Ogni riga contiene i dettagli di un film che ha un articolo di Wikipedia allegato. Utilizziamo le colonne Title, PlotSummary, Release Year e Origin/Ethnicity.

ds = load_dataset("vishnupriyavr/wiki-movie-plots-with-summaries", split="train")
print(ds)

Connessione al database

A questo punto, iniziamo a configurare Milvus. I passi da seguire sono i seguenti:

1. Creare un database Milvus Lite in un file locale. (Sostituire questo URI con l'indirizzo del server per Milvus Standalone e Milvus Distributed).

client = MilvusClient(uri="./sentence_transformers_example.db")

2. Creare lo schema dei dati. Questo specifica i campi che compongono un elemento, compresa la dimensione dell'incorporazione vettoriale.

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
    FieldSchema(name="year", dtype=DataType.INT64),
    FieldSchema(name="origin", dtype=DataType.VARCHAR, max_length=64),
]

schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
client.create_collection(collection_name=collection_name, schema=schema)

3. Definire l'algoritmo di indicizzazione della ricerca vettoriale. Milvus Lite supporta il tipo di indice FLAT, mentre Milvus Standalone e Milvus Distributed implementano un'ampia varietà di metodi come IVF, HNSW e DiskANN. Per la piccola scala di dati di questa dimostrazione, qualsiasi tipo di indice di ricerca è sufficiente, quindi qui usiamo il più semplice FLAT.

index_params = client.prepare_index_params()
index_params.add_index(field_name="embedding", index_type="FLAT", metric_type="IP")
client.create_index(collection_name, index_params)

Una volta eseguiti questi passaggi, siamo pronti a inserire i dati nella raccolta e a eseguire una ricerca. 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, poiché la ricerca brute force sarà 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.

model = SentenceTransformer("all-MiniLM-L12-v2")

Si esegue un ciclo sulle righe dei dati, si incorpora il campo di riepilogo della trama e si inseriscono le entità nel database dei vettori. In generale, è consigliabile eseguire questa fase su batch di dati per massimizzare il throughput della CPU o della GPU per il modello di embedding, come facciamo qui.

for batch in tqdm(ds.batch(batch_size=512)):
    embeddings = model.encode(batch["PlotSummary"])
    data = [
        {"title": title, "embedding": embedding, "year": year, "origin": origin}
        for title, embedding, year, origin in zip(
            batch["Title"], embeddings, batch["Release Year"], batch["Origin/Ethnicity"]
        )
    ]
    res = client.insert(collection_name=collection_name, data=data)

Esecuzione della ricerca

Dopo aver inserito tutti i dati in Milvus, possiamo iniziare a eseguire le nostre ricerche. In questo esempio, cercheremo i film in base ai riassunti della trama di Wikipedia. Poiché stiamo effettuando una ricerca in batch, il tempo di ricerca è condiviso tra le ricerche dei film. (Riuscite a indovinare quale film avevo in mente di recuperare in base al testo descrittivo della query?).

queries = [
    'A shark terrorizes an LA beach.',
    'An archaeologist searches for ancient artifacts while fighting Nazis.',
    'Teenagers in detention learn about themselves.',
    'A teenager fakes illness to get off school and have adventures with two friends.',
    'A young couple with a kid look after a hotel during winter and the husband goes insane.',
    'Four turtles fight bad guys.'
    ]

# Search the database based on input text
def embed_query(data):
    vectors = model.encode(data)
    return [x for x in vectors]


query_vectors = embed_query(queries)

res = client.search(
    collection_name=collection_name,
    data=query_vectors,
    filter='origin == "American" and year > 1945 and year < 2000',
    anns_field="embedding",
    limit=3,
    output_fields=["title"],
)

for idx, hits in enumerate(res):
    print("Query:", queries[idx])
    print("Results:")
    for hit in hits:
        print(hit["entity"].get("title"), "(", round(hit["distance"], 2), ")")
    print()

I risultati sono:

Query: An archaeologist searches for ancient artifacts while fighting Nazis.
Results:
Love Slaves of the Amazons ( 0.4 )
A Time to Love and a Time to Die ( 0.39 )
The Fifth Element ( 0.39 )

Query: Teenagers in detention learn about themselves.
Results:
The Breakfast Club ( 0.54 )
Up the Academy ( 0.46 )
Fame ( 0.43 )

Query: A teenager fakes illness to get off school and have adventures with two friends.
Results:
Ferris Bueller's Day Off ( 0.48 )
Fever Lake ( 0.47 )
Losin' It ( 0.39 )

Query: A young couple with a kid look after a hotel during winter and the husband goes insane.
Results:
The Shining ( 0.48 )
The Four Seasons ( 0.42 )
Highball ( 0.41 )

Query: Four turtles fight bad guys.
Results:
Teenage Mutant Ninja Turtles II: The Secret of the Ooze ( 0.47 )
Devil May Hare ( 0.43 )
Attack of the Giant Leeches ( 0.42 )