Filmsuche mit Milvus und SentenceTransformers

In diesem Beispiel werden wir eine Wikipedia-Artikelsuche mit Hilfe von Milvus und der SentenceTransformers-Bibliothek durchgehen. Der Datensatz, den wir durchsuchen, ist der Wikipedia-Movie-Plots-Datensatz, der auf Kaggle zu finden ist. Für dieses Beispiel haben wir die Daten in einem öffentlichen Google Drive gehostet.

Fangen wir an.

Voraussetzungen für die Installation

Für dieses Beispiel verwenden wir pymilvus, um uns mit Milvus zu verbinden, sentencetransformers, um Vektoreinbettungen zu erzeugen, und gdown, um den Beispieldatensatz herunterzuladen.

pip install pymilvus sentence-transformers gdown

Abrufen der Daten

Wir werden gdown verwenden, um die ZIP-Datei von Google Drive zu entpacken und sie dann mit der integrierten Bibliothek zipfile zu dekomprimieren.

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

Globale Parameter

Hier finden wir die Hauptargumente, die für die Ausführung mit Ihren eigenen Konten geändert werden müssen. Neben jedem Parameter steht eine Beschreibung seiner Bedeutung.

# 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

Einrichten von Milvus

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

1. Verbinden Sie sich mit der Milvus-Instanz unter Verwendung der angegebenen URI.

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

2. Wenn die Sammlung bereits existiert, löschen Sie sie.

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

3. Erstellen Sie die Sammlung, die die ID, den Titel des Films und die Einbettungen des Handlungstextes enthält.

   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. Erstellen Sie einen Index für die neu erstellte Sammlung und laden Sie sie in den Speicher.

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

Sobald diese Schritte abgeschlossen sind, kann die Sammlung eingefügt und durchsucht werden. 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 eine Brute-Force-Suche auf Daten angewendet wird, die noch indiziert werden müssen.

Einfügen der Daten

In diesem Beispiel wird das miniLM-Modell von SentenceTransformers verwendet, um Einbettungen des Plot-Textes zu erstellen. Dieses Modell liefert 384-Dim-Einbettungen.

In den nächsten Schritten werden wir:

1. Laden der Daten.
2. Einbettung der Plot-Textdaten mit SentenceTransformers.
3. Einfügen der Daten 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()

Ausführen der Suche

Nachdem alle Daten in Milvus eingefügt sind, können wir mit der Suche beginnen. In diesem Beispiel werden wir nach Filmen suchen, die auf dem Plot basieren. Da wir eine Stapelsuche durchführen, wird die Suchzeit auf alle Filme aufgeteilt.

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

Die Ausgabe sollte in etwa so aussehen wie die folgende:

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