Início rápido com o Milvus Lite

Os vectores, o formato de dados de saída dos modelos de Redes Neuronais, podem codificar eficazmente a informação e desempenhar um papel fundamental em aplicações de IA, tais como bases de dados de conhecimento, pesquisa semântica, Retrieval Augmented Generation (RAG) e muito mais.

O Milvus é uma base de dados vetorial de código aberto que se adequa a aplicações de IA de todas as dimensões, desde a execução de um chatbot de demonstração no notebook Jupyter até à criação de uma pesquisa à escala da Web que serve milhares de milhões de utilizadores. Neste guia, vamos mostrar-lhe como configurar o Milvus localmente em minutos e utilizar a biblioteca cliente Python para gerar, armazenar e pesquisar vectores.

Instalar o Milvus

Neste guia, utilizamos o Milvus Lite, uma biblioteca Python incluída em pymilvus que pode ser incorporada na aplicação cliente. O Milvus também suporta a implantação no Docker e no Kubernetes para casos de uso de produção.

Antes de começar, certifique-se de ter o Python 3.8+ disponível no ambiente local. Instale pymilvus que contém a biblioteca do cliente python e o Milvus Lite:

$ pip install -U pymilvus

Configurar a base de dados vetorial

Para criar uma base de dados vetorial Milvus local, basta instanciar um MilvusClient especificando um nome de ficheiro para armazenar todos os dados, tal como "milvus_demo.db".

from pymilvus import MilvusClient

client = MilvusClient("milvus_demo.db")

Criar uma coleção

No Milvus, precisamos de uma coleção para armazenar os vectores e os metadados associados. Pode pensar nela como uma tabela nas bases de dados SQL tradicionais. Ao criar uma coleção, pode definir parâmetros de esquema e de índice para configurar as especificações do vetor, como a dimensionalidade, os tipos de índice e as métricas distantes. Existem também conceitos complexos para otimizar o índice para o desempenho da pesquisa vetorial. Por enquanto, vamos nos concentrar apenas no básico e usar o padrão para tudo o que for possível. No mínimo, só é necessário definir o nome da coleção e a dimensão do campo vetorial da coleção.

if client.has_collection(collection_name="demo_collection"):
    client.drop_collection(collection_name="demo_collection")
client.create_collection(
    collection_name="demo_collection",
    dimension=768,  # The vectors we will use in this demo has 768 dimensions
)

Na configuração acima,

• A chave primária e os campos vectoriais utilizam os seus nomes predefinidos ("id" e "vetor").
• O tipo de métrica (definição da distância vetorial) é definido com o seu valor predefinido(COSINE).
• O campo da chave primária aceita números inteiros e não é incrementado automaticamente (nomeadamente, não utiliza a funcionalidade de auto-id). Em alternativa, pode definir formalmente o esquema da coleção seguindo esta instrução.

Preparar os dados

Neste guia, utilizamos vectores para efetuar pesquisa semântica em texto. Precisamos de gerar vectores para texto descarregando modelos de incorporação. Isso pode ser feito facilmente usando as funções utilitárias da biblioteca pymilvus[model].

Representar texto com vectores

Primeiro, instale a biblioteca de modelos. Este pacote inclui ferramentas de ML essenciais, como o PyTorch. O download do pacote pode demorar algum tempo se o seu ambiente local nunca tiver instalado o PyTorch.

$ pip install "pymilvus[model]"

Gerar embeddings vectoriais com o modelo predefinido. O Milvus espera que os dados sejam inseridos organizados como uma lista de dicionários, em que cada dicionário representa um registo de dados, designado por entidade.

from pymilvus import model

# If connection to https://huggingface.co/ failed, uncomment the following path
# import os
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# This will download a small embedding model "paraphrase-albert-small-v2" (~50MB).
embedding_fn = model.DefaultEmbeddingFunction()

# Text strings to search from.
docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

vectors = embedding_fn.encode_documents(docs)
# The output vector has 768 dimensions, matching the collection that we just created.
print("Dim:", embedding_fn.dim, vectors[0].shape)  # Dim: 768 (768,)

# Each entity has id, vector representation, raw text, and a subject label that we use
# to demo metadata filtering later.
data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(vectors))
]

print("Data has", len(data), "entities, each with fields: ", data[0].keys())
print("Vector dim:", len(data[0]["vector"]))

Dim: 768 (768,)
Data has 3 entities, each with fields:  dict_keys(['id', 'vector', 'text', 'subject'])
Vector dim: 768

[Alternativamente] Usar representação falsa com vectores aleatórios

Se não foi possível descarregar o modelo devido a problemas de rede, como alternativa, pode utilizar vectores aleatórios para representar o texto e ainda assim terminar o exemplo. Tenha apenas em atenção que o resultado da pesquisa não reflectirá a semelhança semântica, uma vez que os vectores são falsos.

import random

# Text strings to search from.
docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]
# Use fake representation with random vectors (768 dimension).
vectors = [[random.uniform(-1, 1) for _ in range(768)] for _ in docs]
data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(vectors))
]

print("Data has", len(data), "entities, each with fields: ", data[0].keys())
print("Vector dim:", len(data[0]["vector"]))

Data has 3 entities, each with fields:  dict_keys(['id', 'vector', 'text', 'subject'])
Vector dim: 768

Inserir dados

Vamos inserir os dados na coleção:

res = client.insert(collection_name="demo_collection", data=data)

print(res)

{'insert_count': 3, 'ids': [0, 1, 2], 'cost': 0}

Pesquisa semântica

Agora podemos fazer pesquisas semânticas representando o texto da consulta de pesquisa como um vetor e realizar pesquisas de semelhança de vectores no Milvus.

Pesquisa vetorial

O Milvus aceita um ou vários pedidos de pesquisa vetorial ao mesmo tempo. O valor da variável query_vectors é uma lista de vectores, em que cada vetor é uma matriz de números float.

query_vectors = embedding_fn.encode_queries(["Who is Alan Turing?"])
# If you don't have the embedding function you can use a fake vector to finish the demo:
# query_vectors = [ [ random.uniform(-1, 1) for _ in range(768) ] ]

res = client.search(
    collection_name="demo_collection",  # target collection
    data=query_vectors,  # query vectors
    limit=2,  # number of returned entities
    output_fields=["text", "subject"],  # specifies fields to be returned
)

print(res)

data: ["[{'id': 2, 'distance': 0.5859944820404053, 'entity': {'text': 'Born in Maida Vale, London, Turing was raised in southern England.', 'subject': 'history'}}, {'id': 1, 'distance': 0.5118255615234375, 'entity': {'text': 'Alan Turing was the first person to conduct substantial research in AI.', 'subject': 'history'}}]"] , extra_info: {'cost': 0}

A saída é uma lista de resultados, cada um mapeado para uma consulta de pesquisa vetorial. Cada consulta contém uma lista de resultados, onde cada resultado contém a chave primária da entidade, a distância para o vetor de consulta e os detalhes da entidade com o especificado output_fields.

Pesquisa vetorial com filtragem de metadados

Também é possível efetuar uma pesquisa vetorial tendo em conta os valores dos metadados (chamados campos "escalares" no Milvus, uma vez que escalares se referem a dados não vectoriais). Isto é feito com uma expressão de filtro que especifica determinados critérios. Vejamos como pesquisar e filtrar com o campo subject no exemplo seguinte.

# Insert more docs in another subject.
docs = [
    "Machine learning has been used for drug design.",
    "Computational synthesis with AI algorithms predicts molecular properties.",
    "DDR1 is involved in cancers and fibrosis.",
]
vectors = embedding_fn.encode_documents(docs)
data = [
    {"id": 3 + i, "vector": vectors[i], "text": docs[i], "subject": "biology"}
    for i in range(len(vectors))
]

client.insert(collection_name="demo_collection", data=data)

# This will exclude any text in "history" subject despite close to the query vector.
res = client.search(
    collection_name="demo_collection",
    data=embedding_fn.encode_queries(["tell me AI related information"]),
    filter="subject == 'biology'",
    limit=2,
    output_fields=["text", "subject"],
)

print(res)

data: ["[{'id': 4, 'distance': 0.27030569314956665, 'entity': {'text': 'Computational synthesis with AI algorithms predicts molecular properties.', 'subject': 'biology'}}, {'id': 3, 'distance': 0.16425910592079163, 'entity': {'text': 'Machine learning has been used for drug design.', 'subject': 'biology'}}]"] , extra_info: {'cost': 0}

Por defeito, os campos escalares não são indexados. Se precisar de efetuar uma pesquisa filtrada de metadados num grande conjunto de dados, pode considerar a utilização de um esquema fixo e também ativar o índice para melhorar o desempenho da pesquisa.

Para além da pesquisa vetorial, também pode efetuar outros tipos de pesquisa:

Consulta

Uma query() é uma operação que recupera todas as entidades que correspondem a um critério, tal como uma expressão de filtro ou a correspondência de alguns ids.

Por exemplo, recuperar todas as entidades cujo campo escalar tem um determinado valor:

res = client.query(
    collection_name="demo_collection",
    filter="subject == 'history'",
    output_fields=["text", "subject"],
)

Recuperar diretamente entidades por chave primária:

res = client.query(
    collection_name="demo_collection",
    ids=[0, 2],
    output_fields=["vector", "text", "subject"],
)

Eliminar entidades

Se pretender purgar dados, pode eliminar entidades especificando a chave primária ou eliminar todas as entidades que correspondam a uma determinada expressão de filtro.

# Delete entities by primary key
res = client.delete(collection_name="demo_collection", ids=[0, 2])

print(res)

# Delete entities by a filter expression
res = client.delete(
    collection_name="demo_collection",
    filter="subject == 'biology'",
)

print(res)

[0, 2]
[3, 4, 5]

Carregar dados existentes

Uma vez que todos os dados do Milvus Lite são armazenados num ficheiro local, pode carregar todos os dados para a memória mesmo após o programa terminar, criando um MilvusClient com o ficheiro existente. Por exemplo, isto irá recuperar as colecções do ficheiro "milvus_demo.db" e continuar a escrever dados no mesmo.

from pymilvus import MilvusClient

client = MilvusClient("milvus_demo.db")

Eliminar a coleção

Se pretender eliminar todos os dados de uma coleção, pode eliminar a coleção com

# Drop collection
client.drop_collection(collection_name="demo_collection")

Saiba mais

O Milvus Lite é ótimo para começar com um programa python local. Se você tem dados em grande escala ou gostaria de usar o Milvus em produção, pode aprender sobre a implantação do Milvus no Docker e no Kubernetes. Todos os modos de implantação do Milvus compartilham a mesma API, portanto, seu código do lado do cliente não precisa mudar muito se mudar para outro modo de implantação. Basta especificar o URI e o Token de um servidor Milvus implantado em qualquer lugar:

client = MilvusClient(uri="http://localhost:19530", token="root:Milvus")

Milvus fornece API REST e gRPC, com bibliotecas de clientes em linguagens como Python, Java, Go, C# e Node.js.