Búsqueda de películas utilizando Milvus y SentenceTransformers

En este ejemplo, buscaremos resúmenes de argumentos de películas utilizando Milvus y la biblioteca SentenceTransformers. El conjunto de datos que utilizaremos es Wikipedia Movie Plots with Summaries alojado en HuggingFace.

¡Vamos a empezar!

Bibliotecas necesarias

Para este ejemplo, utilizaremos pymilvus para conectarnos y utilizar Milvus, sentence-transformers para generar incrustaciones vectoriales y datasets para descargar el conjunto de datos de ejemplo.

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

Definiremos algunos parámetros globales,

embedding_dim = 384
collection_name = "movie_embeddings"

Descarga y apertura del conjunto de datos

En una sola línea, datasets nos permite descargar y abrir un conjunto de datos. La biblioteca almacenará en caché el conjunto de datos localmente y utilizará esa copia la próxima vez que se ejecute. Cada fila contiene los detalles de una película a la que acompaña un artículo de Wikipedia. Utilizamos las columnas Title, PlotSummary, Release Year, y Origin/Ethnicity.

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

Conexión a la base de datos

Llegados a este punto, vamos a empezar a configurar Milvus. Los pasos son los siguientes:

1. Cree una base de datos Milvus Lite en un archivo local. (Sustituya este URI por la dirección del servidor para Milvus Standalone y Milvus Distributed).

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

2. Cree el esquema de datos. En él se especifican los campos que componen un elemento, incluida la dimensión de la incrustación vectorial.

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. Defina el algoritmo de indexación de búsqueda vectorial. Milvus Lite admite el tipo de índice FLAT, mientras que Milvus Standalone y Milvus Distributed implementan una amplia variedad de métodos como IVF, HNSW y DiskANN. Para la pequeña escala de datos de esta demostración, cualquier tipo de índice de búsqueda es suficiente, por lo que aquí utilizamos el más sencillo, 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 vez realizados estos pasos, estamos listos para insertar datos en la colección y realizar una búsqueda. Cualquier dato añadido será indexado automáticamente y estará disponible para la búsqueda inmediatamente. Si los datos son muy recientes, la búsqueda puede ser más lenta, ya que se utilizará la fuerza bruta en los datos que aún están en proceso de indexación.

Inserción de los datos

En este ejemplo, vamos a utilizar el modelo miniLM de SentenceTransformers para crear incrustaciones del texto del gráfico. Este modelo devuelve incrustaciones de 384 dimensiones.

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

Hacemos un bucle sobre las filas de los datos, incrustamos el campo de resumen del gráfico e insertamos las entidades en la base de datos vectorial. En general, debería realizar este paso sobre lotes de elementos de datos para maximizar el rendimiento de la CPU o GPU para el modelo de incrustación, como hacemos aquí.

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)

Realizar la búsqueda

Con todos los datos insertados en Milvus, podemos empezar a realizar nuestras búsquedas. En este ejemplo, vamos a buscar películas basándonos en los resúmenes argumentales de Wikipedia. Como estamos haciendo una búsqueda por lotes, el tiempo de búsqueda se reparte entre todas las búsquedas de películas. (¿Puedes adivinar qué película tenía en mente recuperar basándome en el texto de descripción de la consulta?)

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

Los resultados son:

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 )