Pesquisa de filmes usando Milvus e SentenceTransformers

Neste exemplo, vamos pesquisar resumos de enredos de filmes usando Milvus e a biblioteca SentenceTransformers. O conjunto de dados que usaremos é Wikipedia Movie Plots with Summaries hospedado no HuggingFace.

Vamos começar!

Bibliotecas necessárias

Para este exemplo, vamos usar pymilvus para nos ligarmos para usar o Milvus, sentence-transformers para gerar embeddings vectoriais, e datasets para descarregar o conjunto de dados de exemplo.

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

Vamos definir alguns parâmetros globais,

embedding_dim = 384
collection_name = "movie_embeddings"

Descarregar e abrir o conjunto de dados

Numa única linha, datasets permite-nos descarregar e abrir um conjunto de dados. A biblioteca armazenará em cache o conjunto de dados localmente e usará essa cópia na próxima vez que for executada. Cada linha contém os detalhes de um filme que tem um artigo da Wikipedia a acompanhar. Utilizamos as colunas Title, PlotSummary, Release Year e Origin/Ethnicity.

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

Ligação à base de dados

Nesta altura, vamos começar a configurar o Milvus. Os passos são os seguintes:

1. Criar uma base de dados Milvus Lite num ficheiro local. (Substitua este URI pelo endereço do servidor do Milvus Standalone e do Milvus Distributed).

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

2. Criar o esquema de dados. Este esquema especifica os campos que compõem um elemento, incluindo a dimensão da incorporação do vetor.

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. Definir o algoritmo de indexação da pesquisa vetorial. Milvus Lite suporta o tipo de índice FLAT, enquanto Milvus Standalone e Milvus Distributed implementam uma grande variedade de métodos como IVF, HNSW e DiskANN. Para a pequena escala de dados nesta demonstração, qualquer tipo de índice de pesquisa é suficiente, por isso usamos aqui o mais simples, 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)

Uma vez concluídos estes passos, estamos prontos a inserir dados na coleção e a efetuar uma pesquisa. Todos os dados adicionados serão indexados automaticamente e ficarão disponíveis para pesquisa imediatamente. Se os dados forem muito recentes, a pesquisa pode ser mais lenta, pois a pesquisa por força bruta será usada em dados que ainda estão em processo de indexação.

Inserindo os dados

Para este exemplo, vamos usar o modelo SentenceTransformers miniLM para criar embeddings do texto do gráfico. Esse modelo retorna embeddings de 384 dimensões.

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

Fazemos um loop nas linhas dos dados, incorporamos o campo de resumo do gráfico e inserimos entidades no banco de dados do vetor. Em geral, você deve executar essa etapa em lotes de itens de dados para maximizar o rendimento da CPU ou GPU para o modelo de incorporação, como fazemos aqui.

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)

Efetuar a pesquisa

Com todos os dados inseridos no Milvus, podemos começar a efetuar as nossas pesquisas. Neste exemplo, vamos procurar filmes com base em resumos de enredos da Wikipédia. Como estamos a fazer uma pesquisa em lote, o tempo de pesquisa é partilhado entre as pesquisas de filmes. (Consegue adivinhar qual o filme que eu tinha em mente para obter com base no texto de descrição da 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()

Os resultados são:

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 )