Carregamento eficiente de dados no Milvus com VectorETL

Neste tutorial, exploraremos como carregar dados de forma eficiente no Milvus usando o VectorETL, uma estrutura ETL leve projetada para bancos de dados vetoriais. O VectorETL simplifica o processo de extração de dados de várias fontes, transformando-os em embeddings vetoriais usando modelos de IA e armazenando-os no Milvus para recuperação rápida e escalável. No final deste tutorial, terá um pipeline ETL funcional que lhe permite integrar e gerir sistemas de pesquisa vetorial com facilidade. Vamos mergulhar de cabeça!

Preparação

Dependência e ambiente

$ pip install --upgrade vector-etl pymilvus

O VectorETL suporta várias fontes de dados, incluindo Amazon S3, Google Cloud Storage, Ficheiro Local, etc. Você pode conferir a lista completa de fontes suportadas aqui. Neste tutorial, vamos nos concentrar no Amazon S3 como um exemplo de fonte de dados.

Vamos carregar documentos do Amazon S3. Portanto, é necessário preparar AWS_ACCESS_KEY_ID e AWS_SECRET_ACCESS_KEY como variáveis de ambiente para acessar com segurança seu bucket S3. Além disso, utilizaremos o modelo de incorporação text-embedding-ada-002 da OpenAI para gerar incorporação para os dados. Também deve preparar a chave api OPENAI_API_KEY como uma variável de ambiente.

import os

os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key-id"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-access-key"

Fluxo de trabalho

Definição da fonte de dados (Amazon S3)

Neste caso, estamos a extrair documentos de um bucket do Amazon S3. O VectorETL permite-nos especificar o nome do bucket, o caminho para os ficheiros e o tipo de dados com que estamos a trabalhar.

source = {
    "source_data_type": "Amazon S3",
    "bucket_name": "my-bucket",
    "key": "path/to/files/",
    "file_type": ".csv",
    "aws_access_key_id": os.environ["AWS_ACCESS_KEY_ID"],
    "aws_secret_access_key": os.environ["AWS_SECRET_ACCESS_KEY"],
}

Configurando o modelo de incorporação (OpenAI)

Assim que tivermos a nossa fonte de dados configurada, precisamos de definir o modelo de incorporação que transformará os nossos dados textuais em incorporações vectoriais. Aqui, usamos o text-embedding-ada-002 do OpenAI neste exemplo.

embedding = {
    "embedding_model": "OpenAI",
    "api_key": os.environ["OPENAI_API_KEY"],
    "model_name": "text-embedding-ada-002",
}

Configurar o Milvus como a base de dados de destino

Precisamos de armazenar os embeddings gerados no Milvus. Aqui, definimos os nossos parâmetros de ligação ao Milvus usando o Milvus Lite.

target = {
    "target_database": "Milvus",
    "host": "./milvus.db",  # os.environ["ZILLIZ_CLOUD_PUBLIC_ENDPOINT"] if using Zilliz Cloud
    "api_key": "",  # os.environ["ZILLIZ_CLOUD_TOKEN"] if using Zilliz Cloud
    "collection_name": "my_collection",
    "vector_dim": 1536,  # 1536 for text-embedding-ada-002
}

Especificar colunas para incorporação

Agora, temos de especificar quais as colunas dos nossos ficheiros CSV que devem ser convertidas em incorporações. Isso garante que apenas os campos de texto relevantes sejam processados, otimizando a eficiência e o armazenamento.

embed_columns = ["col_1", "col_2", "col_3"]

Criando e executando o pipeline VectorETL

Com todas as configurações em vigor, agora inicializamos o pipeline ETL, definimos o fluxo de dados e o executamos.

from vector_etl import create_flow

flow = create_flow()
flow.set_source(source)
flow.set_embedding(embedding)
flow.set_target(target)
flow.set_embed_columns(embed_columns)

# Execute the flow
flow.execute()

Seguindo este tutorial, criamos com sucesso um pipeline ETL de ponta a ponta para mover documentos do Amazon S3 para o Milvus usando o VectorETL. O VectorETL é flexível nas fontes de dados, pelo que pode escolher as fontes de dados que quiser com base nas necessidades específicas da sua aplicação. Com o design modular do VectorETL, pode facilmente estender este pipeline para suportar outras fontes de dados, incorporando modelos, tornando-o uma ferramenta poderosa para IA e fluxos de trabalho de engenharia de dados!