Genera incorporaciones de vectores para datos textuales de forma masiva con DML particionada

En este documento, se explica cómo generar y reabastecer embeddings de vectores de forma masiva para datos textuales (STRING o JSON) que se almacenan en Spanner con SQL y los modelos de embedding de texto de Gemini Enterprise Agent Platform.

Requisitos previos

Debes tener una tabla en tu base de datos de Spanner que contenga datos textuales (STRING o JSON). Para obtener más información sobre la importación de datos, consulta la descripción general de la importación y exportación de Spanner.

Ejemplo de caso de uso

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 embeddings de vectores para la columna description de esta tabla para encontrar artículos similares que se puedan recomendar a los clientes y mejorar su experiencia de compra con la búsqueda de vectores.

Registra un modelo de incorporación

GoogleSQL

Registra un modelo de embedding de texto con el extremo del modelo de Agent Platform 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
)

Reemplaza lo siguiente:

  • MODEL_NAME: Es el nombre del modelo de embedding de texto de Agent Platform.
  • PROJECT: Es el proyecto que aloja el extremo de Agent Platform.
  • LOCATION: Es la ubicación del extremo de Agent Platform.

PostgreSQL

En el dialecto de PostgreSQL, no es necesario registrar el modelo. Pasas el nombre del extremo directamente a la llamada a función spanner.ML_PREDICT_ROW.

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

  • Para mantener el aislamiento de las cuotas, usa un extremo en un proyecto diferente para generar y reabastecer embeddings que el extremo de producción. Reserva el extremo de producción para entregar el tráfico de producción.
  • Asegúrate de que el extremo 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_size para cada región, consulta Obtén embeddings de texto para un fragmento de texto.
  • Define el extremo con una versión de modelo específica (p.ej., @003) en lugar de @latest. Esto se debe a que los vectores de embeddings generados para el mismo fragmento de texto pueden diferir según la versión del modelo que uses, por lo que debes evitar usar diferentes versiones del modelo para generar embeddings en el mismo conjunto de datos. Además, actualizar la versión del modelo en la declaración de definición del modelo no actualiza los embeddings que ya se generaron con este modelo. Una forma de administrar la versión del modelo para los embeddings es crear una columna adicional en la tabla que almacene la versión del modelo.
  • Los modelos de embedding de texto personalizados no son compatibles con las funciones ML.PREDICT de GoogleSQL ni spanner.ML_PREDICT_ROW de PostgreSQL.

Prueba la integración de extremo a extremo del modelo de embeddings

Puedes ejecutar una consulta para probar que el modelo de embeddings se haya configurado correctamente y que se recuperen los embeddings. 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)
);

Reemplaza lo siguiente:

  • MODEL_NAME: Es el nombre del modelo de embedding de texto de 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;

Reemplaza lo siguiente:

  • PROJECT: Es el proyecto que aloja el extremo de Agent Platform.
  • LOCATION: Es la ubicación del extremo de Agent Platform.
  • MODEL_NAME: Es el nombre del modelo de embedding de texto de Agent Platform.

Actualiza la tabla de origen para incluir columnas adicionales en las que se almacenen los embeddings

A continuación, actualiza el esquema de la tabla de origen para agregar una columna de embeddings. Especificar la longitud del vector es opcional para almacenar incorporaciones, se recomienda para el rendimiento de la búsqueda de KNN y es obligatorio para los índices vectoriales. Omite la longitud del vector si la columna almacena vectores de dimensiones variables:

GoogleSQL 

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

Reemplaza lo siguiente:

  • TABLE_NAME: el nombre de la tabla de origen
  • EMBEDDING_COLUMN_NAME: el nombre de la columna en la que deseas agregar los embeddings generados
  • VECTOR_LENGTH: Es la dimensión de salida del modelo de embeddings, como 768 o 3072.

PostgreSQL 

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

Reemplaza lo siguiente:

  • TABLE_NAME: el nombre de la tabla de origen
  • EMBEDDING_COLUMN_NAME: el nombre de la columna en la que deseas agregar los embeddings generados
  • VECTOR_LENGTH: Es la dimensión de salida del modelo de embeddings, como 768 o 3072.

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

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;

Puedes agregar otra columna para administrar la versión del modelo de embeddings.

GoogleSQL 

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT64;

PostgreSQL 

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT8;

Aumenta la cuota de Agent Platform

Es posible que debas aumentar la cuota de la API de Agent Platform para la región que usa el modelo de embedding de texto. Para solicitar un aumento, consulta Aumentos de cuota de Agent Platform.

Reabastece embeddings

Por último, ejecuta la siguiente sentencia UPDATE con DML particionado para generar embeddings para la columna de datos textuales y almacenarlos en tu base de datos. Puedes almacenar la versión del modelo junto con los embeddings. Te recomendamos que ejecutes esta consulta durante un período 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;

Reemplaza lo siguiente:

  • TABLE_NAME: el nombre de la tabla con los datos textuales
  • EMBEDDING_COLUMN_NAME: el nombre de la columna en la que deseas agregar los embeddings generados
  • DATA_COLUMN_NAME: el nombre de la columna con los datos textuales
  • MODEL_NAME: Es el nombre del modelo de embedding de Agent Platform.
  • MAX_ROWS: Es la cantidad máxima de filas por RPC.
  • EMBEDDING_VERSION_COLUMN: la columna que administra la versión del modelo de embedding que se usa para reabastecer tus embeddings
  • MODEL_VERSION: Es la versión del modelo de embedding de texto.
  • FILTER_CONDITION: Es una condición de filtro particionable que deseas aplicar.

Si usas SAFE.ML.PREDICT, se devuelve NULL para las solicitudes con errores. También puedes usar SAFE.ML.PREDICT junto con un filtro WHERE embedding_column IS NULL para volver a ejecutar la consulta sin calcular los embeddings de los campos que ya se calcularon.

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;

Reemplaza lo siguiente:

  • TABLE_NAME: el nombre de la tabla con los datos textuales
  • EMBEDDING_COLUMN_NAME: el nombre de la columna en la que deseas agregar los embeddings generados
  • DATA_COLUMN_NAME: el nombre de la columna con los datos textuales
  • PROJECT: Es el proyecto que aloja el extremo de Agent Platform.
  • LOCATION: Es la ubicación del extremo de Agent Platform.
  • MODEL_NAME: Es el nombre del modelo de embedding de Agent Platform.
  • MODEL_VERSION: Es la versión del modelo de embedding de Agent Platform.
  • MAX_ROWS: Es la cantidad máxima de filas por RPC.
  • EMBEDDING_VERSION_COLUMN: Es la columna que administra la versión del modelo de embedding de texto que se usa para reabastecer tus embeddings.
  • FILTER_CONDITION: Es una condición de filtro particionable que deseas aplicar.

A continuación, se muestra un ejemplo de consulta de reabastecimiento 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 conocer 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. Según la cantidad de incorporaciones que estés reabasteciendo, es posible que debas aumentar este tiempo de espera para asegurarte de que la DML particionada de UPDATE tenga tiempo suficiente para completarse. Para obtener más información, consulta Configura tiempos de espera y reintentos personalizados.

Rendimiento y otras consideraciones

Considera lo siguiente para optimizar el rendimiento cuando se reabastecen los datos de incorporación.

Cantidad de nodos

El DML particionado ejecuta la declaración DML determinada en diferentes particiones en paralelo. En el caso de las instancias con una gran cantidad de nodos, es posible que observes errores de cuota durante la ejecución de DML particionado. Si las solicitudes a la API de Agent Platform se limitan debido a los límites de cuota de la API de Agent Platform, Spanner reintenta estas fallas en el modo de transacción de DML particionada hasta 20 veces. Si observas una tasa alta de errores de cuota en Agent Platform, aumenta la cuota de Agent Platform. También puedes ajustar el paralelismo con la sugerencia a nivel de la sentencia @{pdml_max_parallelism=DESIRED_NUMBER} cuando usas GoogleSQL. En el siguiente ejemplo, se establece el paralelismo en "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;

Tamaño del texto en la columna de datos

El modelo de embedding de Agent Platform tiene límites en la cantidad máxima de tokens para cada entrada de texto. Las diferentes versiones de modelos tienen distintos límites de tokens. Cada solicitud de Agent Platform puede tener varios campos de texto de entrada, pero hay un límite en la cantidad máxima de tokens que se encuentran en una sola solicitud. En el caso de las bases de datos de GoogleSQL, si encuentras un error INVALID_ARGUMENT con el mensaje "La solicitud es demasiado grande", intenta reducir el tamaño del lote para evitar el error. Para ello, puedes configurar default_batch_size o usar la sugerencia de consulta @{remote_udf_max_outstanding_rpcs} cuando registres el modelo.

Cantidad de solicitudes a la API enviadas a Agent Platform

Puedes usar la sugerencia de consulta @{remote_udf_max_outstanding_rpcs} para aumentar o disminuir la cantidad de solicitudes que se envían a Agent Platform desde Spanner. Ten en cuenta que aumentar este límite puede incrementar el uso de CPU y memoria de la instancia de Spanner. En el caso de las bases de datos de GoogleSQL, el uso de esta sugerencia de consulta anula el default_batch_size configurado para tu modelo.

Supervisa el progreso del reabastecimiento

Puedes supervisar la cantidad de solicitudes, la latencia y los bytes de red que se envían a Agent Platform desde Spanner con el panel de estadísticas del sistema.

¿Qué sigue?