Generar incrustaciones de vectores para datos de texto de forma masiva mediante DML particionado

En este documento se explica cómo generar y rellenar inserciones de vectores de forma masiva para datos de texto (STRING o JSON) almacenados en Spanner mediante SQL y los modelos de inserciones de texto de Vertex AI.

Requisitos previos

Tu base de datos de Spanner debe tener una tabla que contenga datos de texto (STRING o JSON). Para obtener más información sobre cómo importar datos, consulta la descripción general de la importación y exportación de Spanner.

Caso práctico de ejemplo

Supongamos que tienes una tabla en Spanner con el siguiente esquema. Esta tabla contiene millones 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)
);

Tu objetivo es generar inserciones vectoriales para la columna description de esta tabla para encontrar artículos similares que recomendar a los clientes y mejorar su experiencia de compra mediante la búsqueda vectorial.

Registrar un modelo de inserción

GoogleSQL

Registra un modelo de inserción de texto con el endpoint del modelo de Vertex AI en tu base de datos de 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
)

Haz los cambios siguientes:

  • MODEL_NAME: el nombre del modelo de inserción de texto de Vertex AI
  • PROJECT: el proyecto que aloja el endpoint de Vertex AI
  • LOCATION: la ubicación del endpoint de Vertex AI

PostgreSQL

En el dialecto PostgreSQL, no es necesario registrar el modelo. El nombre del endpoint se pasa directamente a la llamada de la función spanner.ML_PREDICT_ROW.

Para seguir las prácticas recomendadas, ten en cuenta lo siguiente:

  • Para mantener el aislamiento de las cuotas, usa un endpoint de otro proyecto para generar y rellenar los embeddings en lugar del endpoint de producción. Reserva el endpoint de producción para servir tráfico de producción.
  • Asegúrate de que el endpoint del modelo admita el valor de default_batch_size. Puedes anular el default_batch_size con la sugerencia de consulta @{remote_udf_max_rows_per_rpc=NEW_NUMBER}. Para obtener información sobre el límite de default_batch_sizede cada región, consulta Obtener las inserciones de texto de un fragmento de texto.
  • Define el endpoint con una versión de modelo específica (por ejemplo, @003) en lugar de @latest. Esto se debe a que los vectores de embedding generados para el mismo fragmento de texto pueden variar en función de la versión del modelo que utilices. Por eso, te recomendamos que no uses diferentes versiones del modelo para generar embeddings en el mismo conjunto de datos. Además, si actualiza la versión del modelo en la instrucción de definición del modelo, no se actualizarán las inserciones que ya se hayan generado con este modelo. Una forma de gestionar la versión del modelo de las inserciones es crear una columna adicional en la tabla que almacene la versión del modelo.
  • Los modelos de inserciones de texto personalizados no se admiten con las funciones ML.PREDICT de GoogleSQL y spanner.ML_PREDICT_ROW de PostgreSQL.

Probar la integración completa del modelo de inserciones

Puedes ejecutar una consulta para comprobar que el modelo de incrustaciones se ha configurado correctamente y que se han recuperado las incrustaciones. Por ejemplo, ejecuta la siguiente consulta:

GoogleSQL 

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

Haz los cambios siguientes:

  • MODEL_NAME: el nombre del modelo de inserción de texto de Vertex AI

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;

Haz los cambios siguientes:

  • PROJECT: el proyecto que aloja el endpoint de Vertex AI
  • LOCATION: la ubicación del endpoint de Vertex AI
  • MODEL_NAME: el nombre del modelo de inserción de texto de Vertex AI

Actualiza la tabla de origen para incluir columnas adicionales en las que almacenar las inserciones.

A continuación, actualice el esquema de la tabla de origen para incluir una columna adicional del tipo de datos ARRAY<FLOAT32> para almacenar las inserciones generadas:

GoogleSQL 

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME ARRAY<FLOAT32>;

Haz los cambios siguientes:

  • TABLE_NAME: el nombre de la tabla de origen
  • EMBEDDING_COLUMN_NAME: el nombre de la columna en la que quieres añadir las inserciones generadas

PostgreSQL 

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME real[];

Haz los cambios siguientes:

  • TABLE_NAME: el nombre de la tabla de origen
  • EMBEDDING_COLUMN_NAME: el nombre de la columna en la que quieres añadir las inserciones generadas

Por ejemplo, usando la tabla products, ejecuta lo siguiente:

GoogleSQL 

ALTER TABLE Products
ADD COLUMN desc_embed ARRAY<FLOAT32>;

PostgreSQL 

ALTER TABLE Products
ADD COLUMN desc_embed real[];

Puedes añadir otra columna para gestionar la versión del modelo de inserción.

GoogleSQL 

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT64;

PostgreSQL 

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT8;

Aumentar la cuota de Vertex AI

Es posible que tengas que aumentar la cuota de la API de Vertex AI de la región que usa el modelo de inserciones de texto. Para solicitar un aumento, consulta Aumentos de cuota de Vertex AI.

Para obtener más información, consulta las cuotas y los límites de Vertex AI.

Incrustaciones de backfill

Por último, ejecuta la siguiente declaración UPDATE con DML particionado para generar incrustaciones de la columna de datos de texto y almacenarlas en tu base de datos. Puedes almacenar la versión del modelo junto con las inserciones. Te recomendamos que ejecutes esta consulta durante un periodo de poco tráfico en tu base de datos.

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;

Haz los cambios siguientes:

  • TABLE_NAME: el nombre de la tabla con los datos de texto
  • EMBEDDING_COLUMN_NAME: el nombre de la columna en la que quieres añadir las inserciones generadas
  • DATA_COLUMN_NAME: el nombre de la columna con los datos de texto
  • MODEL_NAME: el nombre del modelo de inserciones de Vertex AI
  • MAX_ROWS: número máximo de filas por RPC
  • EMBEDDING_VERSION_COLUMN: columna que gestiona la versión del modelo de inserción que se usa para rellenar tus inserciones.
  • MODEL_VERSION: la versión del modelo de embedding de texto
  • FILTER_CONDITION: una condición de filtro particionable que quieras aplicar

Si se usa SAFE.ML.PREDICT, se devuelve NULL para las solicitudes fallidas. También puedes usar SAFE.ML.PREDICT junto con un filtro WHERE embedding_column IS NULL para volver a ejecutar tu consulta sin calcular las inserciones de los campos que ya se han calculado.

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;

Haz los cambios siguientes:

  • TABLE_NAME: el nombre de la tabla con los datos de texto
  • EMBEDDING_COLUMN_NAME: el nombre de la columna en la que quieres añadir las inserciones generadas
  • DATA_COLUMN_NAME: el nombre de la columna con los datos de texto
  • PROJECT: el proyecto que aloja el endpoint de Vertex AI
  • LOCATION: la ubicación del endpoint de Vertex AI
  • MODEL_NAME: el nombre del modelo de inserciones de Vertex AI
  • MODEL_VERSION: la versión del modelo de inserción de Vertex AI
  • MAX_ROWS: número máximo de filas por RPC
  • EMBEDDING_VERSION_COLUMN: columna que gestiona la versión del modelo de embedding de texto que se usa para rellenar tus embeddings.
  • FILTER_CONDITION: una condición de filtro particionable que quieras aplicar

Ejemplo de consulta de relleno para la tabla 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 seguir las prácticas recomendadas, ten en cuenta lo siguiente:

  • El tiempo de espera predeterminado de gRPC para la API de Spanner es de una hora. En función de la cantidad de inserciones que estés rellenando, es posible que tengas que aumentar este tiempo de espera para asegurarte de que el DML particionado UPDATE tenga tiempo suficiente para completarse. Para obtener más información, consulta Configurar tiempos de espera y reintentos personalizados.

Rendimiento y otras consideraciones

Tenga en cuenta lo siguiente para optimizar el rendimiento al rellenar datos de inserción.

Número de nodos

DML particionado ejecuta la instrucción DML proporcionada en diferentes particiones en paralelo. En las instancias con un gran número de nodos, es posible que se produzcan errores de cuota durante la ejecución de DML particionado. Si las solicitudes de la API Vertex AI se limitan debido a los límites de cuota de la API Vertex AI, Spanner reintenta estos errores en el modo de transacción DML particionada un máximo de 20 veces. Si observas una tasa alta de errores de cuota en Vertex AI, aumenta la cuota de Vertex AI. También puede ajustar el paralelismo mediante la sugerencia a nivel de instrucción @{pdml_max_parallelism=DESIRED_NUMBER} mientras usa GoogleSQL. En el siguiente ejemplo se asigna el valor "5" al paralelismo:

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;

Tamaño del texto de la columna de datos

El modelo de inserciones de Vertex AI tiene límites en el número máximo de tokens por cada entrada de texto. Las diferentes versiones del modelo tienen límites de tokens distintos. Cada solicitud de Vertex AI puede tener varios campos de texto de entrada, pero hay un límite en el número máximo de tokens que puede haber en una sola solicitud. En las bases de datos de GoogleSQL, si se produce un error INVALID_ARGUMENT con el mensaje "Request is too large" (La solicitud es demasiado grande), prueba a reducir el tamaño del lote para evitar el error. Para ello, puede configurar default_batch_size o usar la sugerencia de consulta @{remote_udf_max_outstanding_rpcs} al registrar el modelo.

Número de solicitudes de API enviadas a Vertex AI

Puedes usar la sugerencia de consulta @{remote_udf_max_outstanding_rpcs} para aumentar o reducir el número de solicitudes enviadas a Vertex AI desde Spanner. Ten en cuenta que, si aumentas este límite, puede aumentar el uso de CPU y memoria de la instancia de Spanner. En las bases de datos de GoogleSQL, al usar esta sugerencia de consulta se anula el default_batch_size configurado para tu modelo.

Monitorizar el progreso del relleno

Puedes monitorizar el número de solicitudes, la latencia y los bytes de red enviados a Vertex AI desde Spanner mediante el panel de control de estadísticas del sistema.

Siguientes pasos