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 de texto (STRING o JSON) que se almacenan en Spanner con SQL y los modelos de embedding de texto de Agent Platform Enterprise de Gemini.

Requisitos previos

Debes tener una tabla en tu base de datos de Spanner que contenga datos de texto (STRING o JSON). Para obtener más información sobre la importación de datos, consulta la descripción general de 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.

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

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 elementos similares que se recomienden a los clientes y mejorar su experiencia de compra con la búsqueda vectorial.

Registra un modelo de embedding

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
)

Para conocer las prácticas recomendadas, considera 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 tráfico de producción.
• Asegúrate de que el extremo del modelo admita el valor de default_batch_size. Puedes anular default_batch_size con la sugerencia de consulta @{remote_udf_max_rows_per_rpc=NEW_NUMBER}. Para obtener información sobre el default_batch_size límite 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 embedding 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 instrucció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 ajustados de forma personalizada no son compatibles con las funciones ML.PREDICT de GoogleSQL y 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 configuró correctamente y que se recuperan los embeddings. Por ejemplo, ejecuta la siguiente consulta:

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

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;

Actualiza la tabla de origen para incluir columnas adicionales para almacenar los embeddings

A continuación, actualiza el esquema de la tabla de origen para incluir una columna adicional del tipo de datos ARRAY<FLOAT32> para almacenar los embeddings generados:

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME ARRAY<FLOAT32>;

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME real[];

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

ALTER TABLE Products
ADD COLUMN desc_embed ARRAY<FLOAT32>;

ALTER TABLE Products
ADD COLUMN desc_embed real[];

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

ALTER TABLE Products
ADD COLUMN desc_embed_model_version INT64;

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 los embeddings

Por último, ejecuta la siguiente instrucción UPDATE con DML particionado para generar embeddings para la columna de datos de texto y almacenar los embeddings 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.

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;

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;

Un ejemplo de consulta de reabastecimiento para la tabla products:

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;

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, considera lo siguiente:

• El tiempo de espera predeterminado de gRPC para la API de Spanner es de una hora. Según la cantidad de embeddings que reabastezcas, es posible que debas 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 Configura tiempos de espera y reintentos personalizados.

Rendimiento y otras consideraciones

Considera lo siguiente para optimizar el rendimiento cuando reabastezcas datos de embeddings.

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 del DML particionado. Si las solicitudes de 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 DML particionado un máximo de 20 veces. Si observas una alta tasa de errores de cuota en Agent Platform, entonces aumenta la cuota de Agent Platform. También puedes ajustar el paralelismo con la sugerencia a nivel de instrucción @{pdml_max_parallelism=DESIRED_NUMBER} mientras usas GoogleSQL. En el siguiente ejemplo, se establece el paralelismo en “5”:

@{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 embeddings de Agent Platform tiene límites en la cantidad máxima de tokens para cada entrada de texto. Las diferentes versiones del modelo tienen diferentes 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 un mensaje “Request is too large”, 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 enviadas a Agent Platform desde Spanner. Ten en cuenta que aumentar este límite puede aumentar 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 enviados a Agent Platform desde Spanner con el panel de estadísticas del sistema.

¿Qué sigue?

• Obtén información para realizar una búsqueda de vectores de similitud encontrando los K vecinos más cercanos.
• Obtén más información sobre el aprendizaje automático y los embeddings en nuestro curso intensivo sobre embeddings.
• Obtén más información sobre los modelos de embedding de texto de Agent Platform.