Gerar embeddings de vetor para dados textuais em massa usando DML particionada

Este documento explica como gerar e preencher embeddings de vetor em massa para dados textuais (STRING ou JSON) armazenados no Spanner usando SQL e o modelo textembedding-gecko da Vertex AI.

Pré-requisitos

Você precisa ter uma tabela no seu banco de dados do Spanner que contenha dados textuais (STRING ou JSON). Para mais informações sobre a importação de dados, consulte a Visão geral de importação e exportação do Spanner.

Exemplo de caso de uso:

Suponha que você tenha uma tabela no Spanner com o seguinte esquema. Essa tabela contém milhões de registros.

GoogleSQL 

CREATE TABLE Products (
  product_id INT64 NOT NULL,
  name STRING(MAX),
  description STRING(MAX)
) PRIMARY KEY(product_id);

PostgreSQL 

CREATE TABLE Products (
  product_id INT8 NOT NULL,
  name TEXT,
  description TEXT,
  PRIMARY KEY(product_id)
);

Seu objetivo é gerar embeddings vetoriais para a coluna description nesta tabela para encontrar itens semelhantes a serem recomendados aos clientes para melhorar a experiência de compra usando a pesquisa vetorial.

Registrar um modelo de embedding

GoogleSQL

Registre um modelo de incorporação com o endpoint textembedding-gecko da Vertex AI no seu banco de dados do Spanner:

CREATE MODEL MODEL_NAME
INPUT(
  content STRING(MAX)
)
OUTPUT(
  embeddings STRUCT<values ARRAY<FLOAT32>>
)
REMOTE OPTIONS(
    endpoint = '//aiplatform.googleapis.com/projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko$MODEL_VERSION',
  default_batch_size = 5
)

Substitua:

  • MODEL_NAME: o nome do modelo de embedding
  • PROJECT: o projeto que hospeda o endpoint da Vertex AI
  • LOCATION: o local do endpoint da Vertex AI
  • MODEL_VERSION: a versão do modelo de embedding textembedding-gecko

PostgreSQL

No dialeto do PostgreSQL, não é necessário registrar o modelo. Transmita o nome do endpoint diretamente para a chamada de função spanner.ML_PREDICT_ROW.

Para conferir as práticas recomendadas, considere o seguinte:

  • Para manter o isolamento das cotas, use um endpoint em um projeto diferente para gerar e preencher as incorporações do que o endpoint de produção. Reserve o endpoint de produção para veicular o tráfego de produção.
  • Verifique se o endpoint do modelo aceita o valor de default_batch_size. É possível substituir o default_batch_size com a sugestão de consulta @{remote_udf_max_rows_per_rpc=NEW_NUMBER}. Para saber mais sobre o limite de default_batch_size para cada região, consulte Usar embeddings de texto para um snippet de texto.
  • Defina o endpoint com uma versão específica do modelo (por exemplo, @003) em vez de @latest. Isso ocorre porque os vetores de inserção gerados para o mesmo texto podem ser diferentes dependendo da versão do modelo usada. Por isso, é importante evitar o uso de versões diferentes do modelo para gerar embeddings no mesmo conjunto de dados. Além disso, atualizar a versão do modelo na declaração de definição não atualiza as embeddings que já foram geradas com esse modelo. Uma maneira de gerenciar a versão do modelo para embeddings é criar outra coluna na tabela que armazena a versão do modelo.
  • Os modelos textembedding-gecko ajustados personalizados não são compatíveis com a função ML.PREDICT do GoogleSQL e a função spanner.ML_PREDICT_ROW do PostgreSQL.

Testar a integração completa do modelo de embeddings

É possível executar uma consulta para testar se o modelo de inserção está configurado corretamente e se as embeddings são recuperadas. Por exemplo, execute a seguinte consulta:

GoogleSQL 

SELECT embeddings.values
FROM SAFE.ML.PREDICT(
  MODEL MODEL_NAME,
  (SELECT description AS content FROM products LIMIT 10)
);

Substitua:

  • MODEL_NAME: o nome do modelo de embedding

PostgreSQL 

SELECT spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko$MODEL_VERSION',
    JSONB_BUILD_OBJECT('instances', JSONB_BUILD_ARRAY(JSONB_BUILD_OBJECT('content', description))))
FROM Products
LIMIT 10;

Substitua:

  • PROJECT: o projeto que hospeda o endpoint da Vertex AI
  • LOCATION: o local do endpoint da Vertex AI
  • MODEL_VERSION: a versão do modelo de embedding textembedding-gecko

Atualizar a tabela de origem para incluir outras colunas para armazenar os embeddings

Em seguida, atualize o esquema da tabela de origem para incluir outra coluna do tipo de dados ARRAY<FLOAT32> e armazenar as incorporações geradas:

GoogleSQL 

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME ARRAY<FLOAT32>;

Substitua:

  • TABLE_NAME: o nome da tabela de origem
  • EMBEDDING_COLUMN_NAME: o nome da coluna em que você quer adicionar embeddings gerados

PostgreSQL 

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME real[];

Substitua:

  • TABLE_NAME: o nome da tabela de origem
  • EMBEDDING_COLUMN_NAME: o nome da coluna em que você quer adicionar embeddings gerados

Por exemplo, usando o exemplo da tabela products, execute:

GoogleSQL 

ALTER TABLE Products
ADD COLUMN desc_embed ARRAY<FLOAT32>;

PostgreSQL 

ALTER TABLE Products
ADD COLUMN desc_embed real[];

É possível adicionar outra coluna para gerenciar a versão do modelo de incorporação.

GoogleSQL 

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT64;

PostgreSQL 

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT8;

Aumentar a cota da Vertex AI

Talvez seja necessário aumentar a cota da API Vertex AI para textembedding-gecko na região que usa o modelo. Para solicitar um aumento, consulte Aumentos de cota da Vertex AI.

Para mais informações, consulte Cotas e limites da Vertex AI.

Preencher embeddings

Por fim, execute a instrução UPDATE a seguir usando a DML particionada para gerar embeddings para a coluna de dados textuais e armazená-las no banco de dados. Você pode armazenar a versão do modelo com os embeddings. Recomendamos que você execute essa consulta durante uma janela de baixo tráfego no seu banco de dados.

GoogleSQL 

UPDATE TABLE_NAME
SET 
  TABLE_NAME.EMBEDDING_COLUMN_NAME = (
    SELECT embeddings.values
    FROM SAFE.ML.PREDICT(
      MODEL MODEL_NAME,
      (SELECT TABLE_NAME.DATA_COLUMN_NAME AS content)
    ) @{remote_udf_max_rows_per_rpc=MAX_ROWS}
  ),
  TABLE_NAME.EMBEDDING_VERSION_COLUMN = MODEL_VERSION
WHERE FILTER_CONDITION;

Substitua:

  • TABLE_NAME: o nome da tabela com os dados textuais
  • EMBEDDING_COLUMN_NAME: o nome da coluna em que você quer adicionar embeddings gerados
  • DATA_COLUMN_NAME: o nome da coluna com os dados textuais
  • MODEL_NAME: o nome do modelo de embedding
  • MAX_ROWS: o número máximo de linhas por RPC
  • EMBEDDING_VERSION_COLUMN: a coluna que gerencia a versão do modelo de embedding textembedding-gecko usado para preencher os embeddings
  • MODEL_VERSION: a versão do modelo de embedding textembedding-gecko
  • FILTER_CONDITION: uma condição de filtro particionável que você quer aplicar

O uso de SAFE.ML.PREDICT retorna NULL para solicitações com falha. Também é possível usar SAFE.ML.PREDICT em combinação com um filtro WHERE embedding_column IS NULL para executar a consulta novamente sem calcular as representações para os campos que já foram computados.

PostgreSQL 

UPDATE TABLE_NAME
SET 
  EMBEDDING_COLUMN_NAME = spanner.FLOAT32_ARRAY(spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko$MODEL_VERSION', 
    JSONB_BUILD_OBJECT('instances', JSONB_BUILD_ARRAY(JSONB_BUILD_OBJECT('content', DATA_COLUMN_NAME)))
  ) /*@ remote_udf_max_rows_per_rpc=MAX_ROWS */ ->'predictions'->0->'embeddings'->'values'),
  EMBEDDING_VERSION_COLUMN = MODEL_VERSION
WHERE FILTER_CONDITION;

Substitua:

  • TABLE_NAME: o nome da tabela com os dados textuais
  • EMBEDDING_COLUMN_NAME: o nome da coluna em que você quer adicionar embeddings gerados
  • DATA_COLUMN_NAME: o nome da coluna com os dados textuais
  • PROJECT: o projeto que hospeda o endpoint da Vertex AI
  • LOCATION: o local do endpoint da Vertex AI
  • MODEL_VERSION: a versão do modelo de embedding textembedding-gecko
  • MAX_ROWS: o número máximo de linhas por RPC
  • EMBEDDING_VERSION_COLUMN: a coluna que gerencia a versão do modelo de embedding textembedding-gecko usado para preencher os embeddings
  • FILTER_CONDITION: uma condição de filtro particionável que você quer aplicar

Exemplo de consulta de preenchimento para a tabela products:

GoogleSQL 

UPDATE products
SET
  products.desc_embed = (
    SELECT embeddings.values
    FROM SAFE.ML.PREDICT(
      MODEL gecko_model,
      (SELECT products.description AS content)
    ) @{remote_udf_max_rows_per_rpc=200}
  ),
  products.desc_embed_model_version = 3
WHERE products.desc_embed IS NULL;

PostgreSQL 

UPDATE products
SET
  desc_embed = spanner.FLOAT32_ARRAY(spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko@003', 
    JSONB_BUILD_OBJECT('instances', JSONB_BUILD_ARRAY(JSONB_BUILD_OBJECT('content', description)))
  ) /*@ remote_udf_max_rows_per_rpc=200 */ ->'predictions'->0->'embeddings'->'values'),
  desc_embed_model_version = 3
WHERE desc_embed IS NULL;

Para conferir as práticas recomendadas, considere o seguinte:

  • O tempo limite padrão do gRPC para a API Spanner é de uma hora. Dependendo da quantidade de incorporações que você está preenchendo, talvez seja necessário aumentar esse tempo limite para garantir que a DML particionada UPDATE tenha tempo suficiente para ser concluída. Para mais informações, consulte Configurar tempos limite e novas tentativas personalizados.

Desempenho e outras considerações

Considere o seguinte para otimizar o desempenho ao preencher dados de incorporação.

Número de nós

A DML particionada executa a instrução DML especifica em diferentes partições em paralelo. Em instâncias com um grande número de nós, você pode observar erros de cota durante a execução da DML particionada. Se as solicitações da API da Vertex AI forem limitadas devido aos limites de cota da API da Vertex AI, o Spanner vai tentar essas falhas novamente no modo de transação DML particionado por um máximo de 20 vezes. Se você notar uma taxa alta de erros de cota na Vertex AI, aumente a cota dela. Também é possível ajustar o paralelismo usando a sugestão no nível da instrução @{pdml_max_parallelism=DESIRED_NUMBER} ao usar o GoogleSql. O exemplo a seguir define o paralelismo como '5':

GoogleSQL 

@{pdml_max_parallelism=5} UPDATE products
SET products.desc_embed =(
  SELECT embeddings.values
  FROM SAFE.ML.PREDICT(MODEL gecko_model, (
        SELECT products.value AS CONTENT
        )
  )
      @{remote_udf_max_rows_per_rpc=200}
),
products.desc_embed_model_version = 003
WHERE products.desc_embed IS NULL;

Tamanho do texto na coluna de dados

O modelo de incorporação da Vertex AI tem limites no número máximo de tokens para cada entrada de texto. Diferentes versões de modelo têm limites de token diferentes. Cada solicitação da Vertex AI pode ter vários campos de texto de entrada, mas há um limite no número máximo de tokens presentes em uma única solicitação. Para bancos de dados do GoogleSQL, se você encontrar um erro INVALID_ARGUMENT com a mensagem "A solicitação é muito grande", tente reduzir o tamanho do lote para evitar o erro. Para fazer isso, configure default_batch_size ou use a sugestão de consulta @{remote_udf_max_outstanding_rpcs} ao registrar o modelo.

Número de solicitações de API enviadas para a Vertex AI

É possível usar a sugestão de consulta @{remote_udf_max_outstanding_rpcs} para aumentar ou diminuir o número de solicitações enviadas à Vertex AI pelo Spanner. Aumentar esse limite pode aumentar o uso de CPU e memória da instância do Spanner. Para bancos de dados do GoogleSQL, o uso dessa sugestão de consulta substitui o default_batch_size configurado para o modelo.

Monitorar o progresso do preenchimento

É possível monitorar o número de solicitações, a latência e os bytes de rede enviados para a Vertex AI do Spanner usando o painel de insights do sistema.

A seguir