Recherche de films à l'aide de Milvus et de SentenceTransformers

Dans cet exemple, nous allons rechercher des résumés d'intrigues de films en utilisant Milvus et la bibliothèque SentenceTransformers. Le jeu de données que nous utiliserons est Wikipedia Movie Plots with Summaries hébergé sur HuggingFace.

C'est parti !

Bibliothèques requises

Pour cet exemple, nous utiliserons pymilvus pour nous connecter afin d'utiliser Milvus, sentence-transformers pour générer des embeddings vectoriels et datasets pour télécharger le jeu de données de l'exemple.

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

Nous définirons quelques paramètres globaux,

embedding_dim = 384
collection_name = "movie_embeddings"

Téléchargement et ouverture du jeu de données

En une seule ligne, datasets nous permet de télécharger et d'ouvrir un jeu de données. La bibliothèque mettra le jeu de données en cache localement et utilisera cette copie la prochaine fois qu'elle sera exécutée. Chaque ligne contient les détails d'un film accompagné d'un article Wikipedia. Nous utilisons les colonnes Title, PlotSummary, Release Year et Origin/Ethnicity.

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

Connexion à la base de données

A ce stade, nous allons commencer à configurer Milvus. Les étapes sont les suivantes :

1. Créer une base de données Milvus Lite dans un fichier local. (Remplacer cet URI par l'adresse du serveur pour Milvus Standalone et Milvus Distributed).

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

2. Créer le schéma de données. Ce schéma spécifie les champs qui composent un élément, y compris la dimension de l'intégration vectorielle.

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. Définir l'algorithme d'indexation de la recherche vectorielle. Milvus Lite prend en charge le type d'index FLAT, tandis que Milvus Standalone et Milvus Distributed mettent en œuvre une grande variété de méthodes telles que IVF, HNSW et DiskANN. Pour la petite échelle de données de cette démo, n'importe quel type d'index de recherche suffit ; nous utilisons donc ici le plus simple, 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)

Une fois ces étapes franchies, nous sommes prêts à insérer des données dans la collection et à effectuer une recherche. Toutes les données ajoutées seront indexées automatiquement et pourront être recherchées immédiatement. Si les données sont très récentes, la recherche peut être plus lente car la recherche par force brute sera utilisée sur des données qui sont encore en cours d'indexation.

Insérer les données

Pour cet exemple, nous allons utiliser le modèle miniLM SentenceTransformers pour créer des embeddings du texte de l'intrigue. Ce modèle renvoie des embeddings à 384 dimensions.

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

Nous bouclons sur les lignes de données, intégrons le champ de résumé de l'intrigue et insérons les entités dans la base de données vectorielle. En général, il est conseillé d'effectuer cette étape sur des lots d'éléments de données afin de maximiser le débit du CPU ou du GPU pour le modèle d'intégration, comme nous le faisons ici.

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)

Exécution de la recherche

Une fois toutes les données insérées dans Milvus, nous pouvons commencer à effectuer nos recherches. Dans cet exemple, nous allons rechercher des films en nous basant sur les résumés d'intrigue de Wikipedia. Comme nous effectuons une recherche par lots, le temps de recherche est partagé entre les recherches de films. (Pouvez-vous deviner quel film j'avais l'intention de rechercher en vous basant sur le texte de description de la requête ?)

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

Les résultats sont les suivants :

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 )