Filmsuche mit Milvus und SentenceTransformers

In diesem Beispiel werden wir mit Milvus und der SentenceTransformers-Bibliothek Zusammenfassungen von Filmplots suchen. Der Datensatz, den wir verwenden werden, ist Wikipedia Movie Plots with Summaries, der auf HuggingFace gehostet wird.

Los geht's!

Erforderliche Bibliotheken

Für dieses Beispiel verwenden wir pymilvus, um eine Verbindung zu Milvus herzustellen, sentence-transformers, um Vektoreinbettungen zu erzeugen, und datasets, um den Beispieldatensatz herunterzuladen.

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

Wir werden einige globale Parameter definieren,

embedding_dim = 384
collection_name = "movie_embeddings"

Herunterladen und Öffnen des Datensatzes

In einer einzigen Zeile können wir mit datasets einen Datensatz herunterladen und öffnen. Die Bibliothek wird den Datensatz lokal zwischenspeichern und beim nächsten Aufruf diese Kopie verwenden. Jede Zeile enthält die Details eines Films, zu dem ein Wikipedia-Artikel existiert. Wir verwenden die Spalten Title, PlotSummary, Release Year und Origin/Ethnicity.

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

Verbinden mit der Datenbank

An dieser Stelle beginnen wir mit der Einrichtung von Milvus. Die Schritte sind wie folgt:

1. Erstellen Sie eine Milvus-Lite-Datenbank in einer lokalen Datei. (Ersetzen Sie diese URI durch die Serveradresse für Milvus Standalone und Milvus Distributed).

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

2. Erstellen Sie das Datenschema. Dieses gibt die Felder an, aus denen ein Element besteht, einschließlich der Dimension der Vektoreinbettung.

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. Definieren Sie den Indexierungsalgorithmus für die Vektorsuche. Milvus Lite unterstützt den Index-Typ FLAT, während Milvus Standalone und Milvus Distributed eine Vielzahl von Methoden wie IVF, HNSW und DiskANN implementieren. Für den kleinen Umfang der Daten in dieser Demo reicht jeder Suchindex-Typ aus, so dass wir hier den einfachsten Typ FLAT verwenden.

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)

Sobald diese Schritte abgeschlossen sind, können wir Daten in die Sammlung einfügen und eine Suche durchführen. Alle hinzugefügten Daten werden automatisch indiziert und sind sofort für die Suche verfügbar. Wenn die Daten sehr frisch sind, kann die Suche langsamer sein, da die Brute-Force-Suche auf Daten angewendet wird, die noch indiziert werden müssen.

Einfügen der Daten

In diesem Beispiel wird das miniLM-Modell SentenceTransformers verwendet, um Einbettungen des Plottextes zu erstellen. Dieses Modell liefert 384-dimensionale Einbettungen.

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

Wir durchlaufen eine Schleife über die Datenzeilen, betten das Feld mit der Zusammenfassung der Darstellung ein und fügen die Entitäten in die Vektordatenbank ein. Im Allgemeinen sollten Sie diesen Schritt über Stapel von Datenelementen ausführen, um den CPU- oder GPU-Durchsatz für das Einbettungsmodell zu maximieren, wie wir es hier tun.

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)

Ausführen der Suche

Nachdem alle Daten in Milvus eingegeben wurden, können wir mit der Suche beginnen. In diesem Beispiel suchen wir nach Filmen auf der Grundlage von Handlungszusammenfassungen aus Wikipedia. Da wir eine Stapelsuche durchführen, wird die Suchzeit auf alle Filme aufgeteilt. (Können Sie anhand des Beschreibungstextes der Abfrage erraten, welchen Film ich im Sinn hatte?)

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

Die Ergebnisse sind:

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 )