Pesquisa de filmes usando Milvus e SentenceTransformers

Neste exemplo, vamos analisar uma pesquisa de artigos da Wikipédia usando Milvus e a biblioteca SentenceTransformers. O conjunto de dados que estamos a pesquisar é o Wikipedia-Movie-Plots Dataset encontrado no Kaggle. Para este exemplo, colocámos os dados num Google Drive público.

Vamos começar.

Requisitos de instalação

Para este exemplo, vamos usar pymilvus para nos ligarmos para usar o Milvus, sentencetransformers para gerar embeddings vectoriais e gdown para descarregar o conjunto de dados de exemplo.

pip install pymilvus sentence-transformers gdown

Obter os dados

Vamos utilizar gdown para obter o zip do Google Drive e depois descomprimi-lo com a biblioteca zipfile incorporada.

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

Parâmetros globais

Aqui podemos encontrar os principais argumentos que precisam de ser modificados para serem executados com as suas próprias contas. Ao lado de cada um há uma descrição do que se trata.

# 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

Configurar o Milvus

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

1. Ligue-se à instância do Milvus utilizando o URI fornecido.

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

2. Se a coleção já existir, elimine-a.

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

3. Crie a coleção que contém o id, o título do filme e os embeddings do texto do enredo.

   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. Crie um índice na coleção recém-criada e carregue-a para a memória.

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

Quando estes passos estiverem concluídos, a coleção estará pronta para ser inserida e pesquisada. Todos os dados adicionados serão indexados automaticamente e ficarão imediatamente disponíveis para pesquisa. Se os dados forem muito recentes, a pesquisa pode ser mais lenta, uma vez que a pesquisa por força bruta será utilizada em dados que ainda estão a ser indexados.

Inserir os dados

Para este exemplo, vamos utilizar o modelo SentenceTransformers miniLM para criar embeddings do texto do gráfico. Esse modelo retorna embeddings de 384 dígitos.

Nas próximas etapas, nós vamos:

1. Carregando os dados.
2. Incorporar os dados do texto do enredo usando SentenceTransformers.
3. Inserir os dados no 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()

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 no enredo. Como estamos a fazer uma pesquisa em lote, o tempo de pesquisa é partilhado entre as pesquisas de filmes.

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

O resultado deve ser semelhante ao seguinte:

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