Criar um pipeline de embedding de vetor em tempo real para o AlloyDB com o Dataflow

Este documento mostra como criar um pipeline de extração, transformação e carregamento (ETL) do AlloyDB para PostgreSQL usando o Dataflow. Google Cloud O Dataflow é um serviço Google Cloud totalmente gerenciado para desenvolver e executar pipelines de processamento de dados.

Você pode usar as instruções no documento, que são baseadas no Colab Ingestão de embeddings de vetores com Apache Beam e AlloyDB, que usa Python para criar o pipeline de ingestão basic_ingestion_pipeline.py. Alguns dos casos de uso em que você pode aplicar as informações deste documento são pesquisa semântica ou geração aumentada de recuperação (RAG).

Estas instruções descrevem os seguintes componentes do pipeline do Dataflow:

  • Como configurar uma conexão do AlloyDB e do Dataflow
  • Gerar embeddings no AlloyDB para PostgreSQL usando o manipulador VertexAITextEmbeddings do Apache Beam e o modelo de embedding de texto da Vertex AI
  • Criar um pipeline de streaming no Dataflow

Antes de começar

Antes de criar o pipeline do Dataflow usando o Colab, atenda aos seguintes pré-requisitos:

Configurar a instância do AlloyDB para PostgreSQL e os componentes do pipeline

Primeiro, configure o pipeline para se conectar a uma instância do AlloyDB para PostgreSQL. Essa configuração inclui a definição do ID do projeto Google Cloud , do URI da instância do AlloyDB para PostgreSQL, do usuário e da senha para se conectar usando o conector de linguagem do AlloyDB. Para mais informações sobre como configurar a conexão, consulte Configuração do banco de dados.

Os módulos do Apache Beam específicos da geração aumentada de recuperação (RAG) fornecem classes para as seguintes tarefas:

  • Ingerir dados do AlloyDB para PostgreSQL
  • Gerando embeddings
  • Gravar esses embeddings de volta no AlloyDB para PostgreSQL

Importe as classes necessárias para o código do pipeline antes de criar a lógica dele. Para mais informações sobre componentes de pipeline, consulte Importar componentes de pipeline.

Criar dados de amostra

O Colab Ingestão de embeddings de vetor com Apache Beam e AlloyDB fornece dados de amostra products_data para executar o pipeline. O pipeline usa esses dados de amostra como entrada, junto com o modelo de embedding, para gerar embeddings.

Para mais informações, consulte Criar dados de amostra.

Criar uma tabela para armazenar embeddings

O pipeline armazena os embeddings gerados na tabela default_dataflow_product_embeddings. Para mais informações sobre como criar o esquema da tabela, consulte Criar tabela com esquema padrão.

Opcional: preparar dados para ingestão de embeddings

Com base no seu conjunto de dados, você pode dividir os dados em metadados e texto que o modelo de embedding precisa converter em embeddings. As classes MLTransform() e VectorDatabaseWriteTransform() processam dados de entrada em um tamanho compatível com o modelo de incorporação. Inclua os metadados e formate os dados de entrada de acordo com as especificações do modelo de incorporação que você está usando.

Para mais informações sobre como preparar dados, consulte Mapear dados de produtos para partes.

Configurar o manipulador de embeddings para gerar embeddings

A classe VertexAITextEmbeddings() define o modelo de embedding de texto que cria embeddings de vetores. Esse modelo de embedding converte os dados em partes em embeddings.

Para mais informações, consulte Configurar o manipulador de incorporação.

Também é possível usar um modelo pré-treinado criado com o framework SentenceTransformers do Huggingface para gerar embeddings de vetor. Para mais informações, consulte Gerar embeddings com o HuggingFace.

Criar um pipeline de ingestão

O pipeline basic_ingestion_pipeline.py, fornecido no Colab Ingestão de embeddings de vetores com Apache Beam e AlloyDB, incorpora as configurações das seções anteriores, incluindo a configuração do AlloyDB para PostgreSQL, o carregamento de dados no AlloyDB para PostgreSQL, o chunking de dados opcional e a configuração do manipulador de embeddings.

O pipeline de ingestão faz o seguinte:

  • Cria tabelas de dados de produtos
  • Converte dados em partes
  • Gera embeddings
  • Grava os embeddings convertidos na tabela products_data do AlloyDB para PostgreSQL.

É possível executar esse pipeline usando um executor local direto ou um executor baseado na nuvem, como o Dataflow.

Para mais informações sobre como criar o pipeline de ingestão, consulte Salvar nosso pipeline em um arquivo Python.

Execute o pipeline do Dataflow

É possível executar um pipeline do Dataflow na linha de comando. Transmita credenciais, como ID do projeto, detalhes da conexão do AlloyDB para PostgreSQL, local do bucket do Cloud Storage, detalhes do ambiente de execução, informações de rede e o nome do pipeline de ingestão (basic_ingestion_pipeline.py).

No Colab Ingestão de embeddings de vetores com Apache Beam e AlloyDB, a instância do AlloyDB para PostgreSQL e os jobs do Dataflow são executados na mesma rede VPC e sub-rede.

Para mais informações sobre como executar um pipeline no Dataflow, consulte Executar pipeline no Dataflow.

No console do Google Cloud , no painel do Dataflow, é possível conferir gráficos de execução, registros e métricas enquanto o pipeline é executado.

Opcional: execute o pipeline de streaming do Dataflow

Para dados que mudam com frequência, como pesquisas de similaridade ou mecanismos de recomendação, considere criar um pipeline de streaming usando o Dataflow e o Pub/Sub.

Em vez de processar um lote de dados, esse pipeline lê continuamente as mensagens recebidas de um tópico do Pub/Sub, converte as mensagens em partes, gera embeddings usando um modelo especificado (como Hugging Face ou Vertex AI) e atualiza a tabela do AlloyDB para PostgreSQL.

Para mais informações, consulte Streaming de atualizações de incorporações do Pub/Sub.

Verificar embeddings vetoriais no AlloyDB para PostgreSQL

Depois que o pipeline for executado, verifique se ele gravou os embeddings no banco de dados do AlloyDB para PostgreSQL.

Para mais informações, consulte Verificar os embeddings escritos.

A seguir