Gerar embeddings vetoriais para dados textuais em massa usando DML particionada

Este documento explica como gerar e preencher embeddings de vetores em massa. para dados textuais (STRING ou JSON) armazenados em Spanner usando SQL e a Vertex AI textembedding-gecko um modelo de machine learning.

Pré-requisitos

Você precisa 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 vetores para a coluna description neste para encontrar itens semelhantes e recomendar aos clientes para melhorar as compras usando a pesquisa de vetores.

Registrar um modelo de embedding

GoogleSQL

Registrar um modelo de embedding com a Vertex AI Endpoint textembedding-gecko 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/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 de textembedding-gecko.

PostgreSQL

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

Para práticas recomendadas, considere o seguinte:

  • Para manter o isolamento de cotas, use um endpoint em um projeto diferente para gerar e preencher embeddings do que o endpoint de produção. Reserve o de produção para atender ao tráfego de produção.
  • Verifique se o endpoint do modelo é compatível com o valor de default_batch_size. Substitua 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 a mesma parte de texto pode diferir dependendo da versão do modelo que você usa; quais é por isso que é melhor evitar o uso de diferentes versões do modelo para gerar embeddings. no mesmo conjunto de dados. Além disso, atualizar a versão do modelo no modelo de definição não atualiza os embeddings que já foram gerados com esse modelo. Uma maneira de gerenciar a versão do modelo para embeddings é crie outra coluna na tabela que armazene a versão do modelo.
  • Os modelos textembedding-gecko ajustados personalizados não são compatíveis com o GoogleSQL ML.PREDICT e o PostgreSQL função spanner.ML_PREDICT_ROW.

Testar a integração completa do modelo de embedding

É possível executar uma consulta para testar se o modelo de embedding foi configurado corretamente. e 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

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 de textembedding-gecko.

Atualizar a tabela de origem para incluir mais colunas e armazenar os embeddings

Em seguida, atualize o esquema da tabela de origem para incluir mais uma coluna do tipo de dados ARRAY<FLOAT32> para armazenar os embeddings gerados:

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 os 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 os 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.

Embeddings de preenchimento

Por fim, execute a seguinte instrução UPDATE usando DML particionada para gerar embeddings para a coluna de dados textuais e armazenar os embeddings em seu banco de dados. É possível armazenar a versão do modelo com os embeddings. Qa recomendamos que você faça essa consulta em uma janela com pouco 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 os 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 seus embeddings
  • MODEL_VERSION: a versão do modelo de embedding de 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. Você também pode usar SAFE.ML.PREDICT em combinação com um WHERE embedding_column IS NULL. filtro 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/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 os 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 de 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 seus embeddings
  • FILTER_CONDITION: uma condição de filtro particionável que você quer aplicar

Um 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 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 estão sendo preenchidos, pode ser necessário aumente esse tempo limite para garantir que a DML particionada por UPDATE tenha tempo suficiente para concluir. 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 a incorporação dados.

Número de nós

A DML particionada executa a instrução DML especificada em partições diferentes parallel. Em instâncias com um grande número de nós, é possível observar erros durante a execução de DML particionada. Se a Vertex AI As solicitações de API são limitadas devido aos limites de cota da API Vertex AI. o Spanner repete essas falhas modo de transação de DML particionada no máximo 20 vezes. Se você observar uma alta taxa de erros de cota no Vertex AI e aumentar a cota da Vertex AI. Você também pode ajustar o paralelismo usando a dica no nível da instrução @{pdml_max_parallelism=DESIRED_NUMBER} ao usar o GoogleSql. O seguinte 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 embedding da Vertex AI tem limites no número máximo de tokens para cada entrada de texto. Diferentes versões de modelo têm diferentes tokens limites. Cada solicitação da Vertex AI pode ter vários textos campos, mas há um limite no número máximo de tokens presentes em uma única solicitação. Para bancos de dados GoogleSQL, se você encontrar uma 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 o 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 à Vertex AI

É possível usar a dica de consulta @{remote_udf_max_outstanding_rpcs} para aumentar ou diminuir o número de solicitações enviadas à Vertex AI de no Spanner. Esteja ciente de que aumentar esse limite pode aumentar e uso da memória da instância do Spanner. Para GoogleSQL bancos de dados, 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 aos Vertex AI do Spanner usando o painel de insights do sistema.

A seguir