Recherche de film avec Milvus et SentenceTransformers

Dans cet exemple, nous allons effectuer une recherche d'articles dans Wikipédia en utilisant Milvus et la bibliothèque SentenceTransformers. Le jeu de données sur lequel nous effectuons notre recherche est le jeu de données Wikipedia-Movie-Plots trouvé sur Kaggle. Pour cet exemple, nous avons réhébergé les données dans un Google Drive public.

Commençons par le début.

Installation des conditions requises

Pour cet exemple, nous allons utiliser pymilvus pour nous connecter à Milvus, sentencetransformers pour générer des embeddings vectoriels et gdown pour télécharger le jeu de données d'exemple.

pip install pymilvus sentence-transformers gdown

Récupérer les données

Nous allons utiliser gdown pour récupérer le fichier zip sur Google Drive et le décompresser avec la bibliothèque intégrée 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")

Paramètres globaux

Nous trouvons ici les principaux arguments qui doivent être modifiés pour fonctionner avec vos propres comptes. A côté de chacun d'entre eux se trouve une description de ce qu'il représente.

# 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

Configuration de Milvus

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

1. Connectez-vous à l'instance Milvus à l'aide de l'URI fourni.

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

2. Si la collection existe déjà, la supprimer.

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

3. Créez la collection qui contient l'identifiant, le titre du film et les embeddings du texte de l'intrigue.

   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. Créez un index sur la collection nouvellement créée et chargez-la en mémoire.

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

Une fois ces étapes réalisées, la collection est prête à être insérée et à faire l'objet de recherches. 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 une recherche par force brute sera utilisée sur les 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 de 384 dim.

Au cours des prochaines étapes, nous allons

1. Charger les données.
2. Intégrer les données du texte de l'intrigue à l'aide de SentenceTransformers.
3. Insérer les données dans 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()

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 fonction de l'intrigue. Comme nous effectuons une recherche par lots, le temps de recherche est partagé entre les recherches de films.

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

La sortie devrait être similaire à ce qui suit :

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