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

Neste documento, explicamos como gerar e preencher embeddings de vetor em massa para dados textuais (STRING ou JSON) armazenados no Spanner usando SQL e os modelos de embedding de texto da Gemini Enterprise Agent Platform.

Pré-requisitos

É preciso ter uma tabela no banco de dados do Spanner que contenha dados textuais (STRING ou JSON). Para mais informações sobre como importar 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 esquema a seguir. 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 de vetor para a coluna description nessa tabela para encontrar itens semelhantes e recomendar aos clientes, melhorando a experiência de compra deles com a busca vetorial.

Registrar um modelo de embedding

GoogleSQL

Registre um modelo de embedding de texto com o endpoint de modelo da Agent Platform no 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/$MODEL_NAME',
  default_batch_size = 5
)

Substitua:

  • MODEL_NAME: o nome do modelo de embedding de texto da Agent Platform.
  • PROJECT: o projeto que hospeda o endpoint da plataforma de agentes.
  • LOCATION: o local do endpoint da Plataforma de Agentes.

PostgreSQL

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

Para as práticas recomendadas, considere o seguinte:

  • Para manter o isolamento das cotas, use um endpoint em um projeto diferente para gerar e preencher embeddings em vez do endpoint de produção. Reserve o endpoint de produção para veicular o tráfego de produção.
  • Verifique se o endpoint do modelo é compatível com o valor de default_batch_size. É possível substituir o default_batch_size pela dica de consulta @{remote_udf_max_rows_per_rpc=NEW_NUMBER}. Para informações sobre o limite de default_batch_size para cada região, consulte Receber embeddings de texto para um snippet de texto.
  • Defina o endpoint com uma versão de modelo específica (por exemplo, @003) em vez de @latest. Isso ocorre porque os vetores de embedding gerados para o mesmo texto podem ser diferentes dependendo da versão do modelo usada. Por isso, evite usar versões de modelo diferentes para gerar embeddings no mesmo conjunto de dados. Além disso, atualizar a versão do modelo na instrução de definição do modelo não atualiza os embeddings que já foram gerados com esse modelo. Uma maneira de gerenciar a versão do modelo para embeddings é criar uma coluna adicional na tabela que armazena a versão do modelo.
  • Os modelos de embedding de texto ajustados de forma personalizada não são compatíveis com as funções ML.PREDICT do GoogleSQL e 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 embedding está configurado corretamente e se os embeddings são recuperados. 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 de texto da Agent Platform.

PostgreSQL 

SELECT spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/$MODEL_NAME',
    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 plataforma de agentes.
  • LOCATION: o local do endpoint da Plataforma de Agentes.
  • MODEL_NAME: o nome do modelo de embedding de texto da Agent Platform.

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

Em seguida, atualize o esquema da tabela de origem para adicionar uma coluna de embedding. Especificar o comprimento do vetor é opcional para armazenar embeddings, recomendado para o desempenho da pesquisa de KNN e obrigatório para índices de vetor. Omita o comprimento do vetor se a coluna armazenar vetores de dimensões variadas:

GoogleSQL 

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME ARRAY<FLOAT32>(vector_length=>VECTOR_LENGTH);

Substitua:

  • TABLE_NAME: o nome da tabela de origem.
  • EMBEDDING_COLUMN_NAME: o nome da coluna em que você quer adicionar os embeddings gerados.
  • VECTOR_LENGTH: a dimensão de saída do modelo de embedding, como 768 ou 3072.

PostgreSQL 

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME real[] VECTOR LENGTH VECTOR_LENGTH;

Substitua:

  • TABLE_NAME: o nome da tabela de origem.
  • EMBEDDING_COLUMN_NAME: o nome da coluna em que você quer adicionar os embeddings gerados.
  • VECTOR_LENGTH: a dimensão de saída do modelo de embedding, como 768 ou 3072.

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

GoogleSQL 

ALTER TABLE Products
ADD COLUMN desc_embed ARRAY<FLOAT32>(vector_length=>768);

PostgreSQL 

ALTER TABLE Products
ADD COLUMN desc_embed real[] VECTOR LENGTH 768;

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

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 Agent Platform

Talvez seja necessário aumentar a cota da API Agent Platform para a região que usa o modelo de embedding de texto. Para solicitar um aumento, consulte Aumentos de cota da plataforma de agentes.

Preencher embeddings

Por fim, execute a seguinte instrução UPDATE usando DML particionada para gerar embeddings para a coluna de dados textuais e armazená-los no banco de dados. É possível armazenar a versão do modelo com os embeddings. Recomendamos que você execute essa consulta durante uma janela de baixo tráfego no 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 os embeddings gerados.
  • DATA_COLUMN_NAME: o nome da coluna com os dados textuais.
  • MODEL_NAME: o nome do modelo de embedding da Agent Platform.
  • 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 usada para preencher os embeddings.
  • MODEL_VERSION: a versão do modelo de embedding de texto.
  • 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 novamente a consulta sem calcular os embeddings dos campos que já foram calculados.

PostgreSQL 

UPDATE TABLE_NAME
SET
  EMBEDDING_COLUMN_NAME = spanner.FLOAT32_ARRAY(spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/$MODEL_NAME',
    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 os embeddings gerados.
  • DATA_COLUMN_NAME: o nome da coluna com os dados textuais.
  • PROJECT: o projeto que hospeda o endpoint da plataforma de agentes.
  • LOCATION: o local do endpoint da Plataforma de Agentes.
  • MODEL_NAME: o nome do modelo de embedding da Agent Platform.
  • MODEL_VERSION: a versão do modelo de embedding da Agent Platform.
  • 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 de texto usada 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 embedding_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/$MODEL_NAME',
    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 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 embeddings 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 embedding.

Número de nós

A DML particionada executa a instrução DML fornecida em diferentes partições em paralelo. Para instâncias com um grande número de nós, é possível observar erros de cota durante a execução da DML particionada. Se as solicitações de API do Agent Platform forem limitadas devido aos limites de cota do Agent Platform, o Spanner vai tentar novamente essas falhas no modo de transação DML particionada por no máximo 20 vezes. Se você observar uma alta taxa de erros de cota na Agent Platform, aumente a cota da Agent Platform. Também é possível ajustar o paralelismo usando a dica 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 embedding_model, (
        SELECT products.value AS CONTENT
        )
  )
      @{remote_udf_max_rows_per_rpc=200}
),
products.desc_embed_model_version = MODEL_VERSION
WHERE products.desc_embed IS NULL;

Tamanho do texto na coluna de dados

O modelo de embedding da Agent Platform tem limites para o número máximo de tokens em cada entrada de texto. Versões de modelo diferentes têm limites de token distintos. Cada solicitação da plataforma de agente pode ter vários campos de texto de entrada, mas há um limite para o 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 uma mensagem "A solicitação é muito grande", tente reduzir o tamanho do lote para evitar o erro. Para isso, configure default_batch_size ou use a dica de consulta @{remote_udf_max_outstanding_rpcs} ao registrar o modelo.

Número de solicitações de API enviadas à plataforma de agentes

Use a dica de consulta @{remote_udf_max_outstanding_rpcs} para aumentar ou diminuir o número de solicitações enviadas do Spanner para a Agent Platform. Esteja ciente de que aumentar esse limite pode aumentar o uso de CPU e o uso da memória da instância do Spanner. Para bancos de dados do GoogleSQL, o uso dessa dica de consulta substitui o default_batch_size configurado para seu 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 plataforma de agentes do Spanner usando o painel de insights do sistema.

A seguir