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

En este documento, se explica cómo generar y reabastecer incorporaciones de vectores de forma masiva para datos textuales (STRING o JSON) que se almacenan en Spanner con SQL y el modelo textembedding-gecko de Vertex AI.

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 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 incorporaciones vectoriales para la columna description en esta tabla para encontrar artículos similares que puedas recomendar a los clientes y mejorar su experiencia de compra mediante la búsqueda vectorial.

Registra un modelo de incorporación

GoogleSQL

Registra un modelo de incorporación con el extremo textembedding-gecko 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/textembedding-gecko$MODEL_VERSION',
  default_batch_size = 5
)

Reemplaza lo siguiente:

  • MODEL_NAME: Es el nombre del modelo de incorporación.
  • PROJECT: Es el proyecto que aloja el extremo de Vertex AI.
  • LOCATION: la ubicación del extremo de Vertex AI
  • MODEL_VERSION: Es la versión del modelo de incorporación de textembedding-gecko.

PostgreSQL

En el dialecto de PostgreSQL, no es necesario registrar el modelo. Pasas el nombre del extremo directamente a la llamada a la 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 incorporaciones que no sean del extremo de producción. Reserva el extremo de producción para publicar 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 incorporaciones 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 incorporación 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 incorporaciones en el mismo conjunto de datos. Además, actualizar la versión del modelo en la sentencia de definición del modelo no actualiza las incorporaciones que ya se generaron con este modelo. Una forma de administrar la versión del modelo para las incorporaciones es crear una columna adicional en la tabla que almacene la versión del modelo.
  • Los modelos textembedding-gecko 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 incorporación se configuró correctamente y que se recuperan las incorporaciones. 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 incorporación.

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;

Reemplaza lo siguiente:

  • PROJECT: Es el proyecto que aloja el extremo de Vertex AI.
  • LOCATION: la ubicación del extremo de Vertex AI
  • MODEL_VERSION: Es la versión del modelo de incorporación de textembedding-gecko.

Actualiza la tabla de origen para incluir columnas adicionales que almacenen las incorporaciones

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

GoogleSQL

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME ARRAY<FLOAT32>;

Reemplaza lo siguiente:

  • TABLE_NAME: Es el nombre de la tabla de origen.
  • EMBEDDING_COLUMN_NAME: El nombre de la columna en la que deseas agregar las incorporaciones generadas

PostgreSQL

ALTER TABLE TABLE_NAME
ADD COLUMN EMBEDDING_COLUMN_NAME real[];

Reemplaza lo siguiente:

  • TABLE_NAME: Es el nombre de la tabla de origen.
  • EMBEDDING_COLUMN_NAME: El nombre de la columna en la que deseas agregar las incorporaciones generadas

Por ejemplo, con el ejemplo de 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 agregar otra columna para administrar la versión del modelo de incorporación.

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 Vertex AI

Es posible que debas aumentar la cuota de la API de Vertex AI para textembedding-gecko en la región que usa el modelo. Para solicitar un aumento, consulta Aumentos de cuota de Vertex AI.

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

Reabastecimiento de incorporaciones

Por último, ejecuta la siguiente sentencia UPDATE con DML particionado para generar incorporaciones para la columna de datos textuales y almacenarlas en tu base de datos. Puedes almacenar la versión del modelo junto con las incorporaciones. 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: Es el nombre de la tabla con los datos textuales.
  • EMBEDDING_COLUMN_NAME: El nombre de la columna en la que deseas agregar las incorporaciones generadas
  • DATA_COLUMN_NAME: Es el nombre de la columna con los datos textuales.
  • MODEL_NAME: Es el nombre del modelo de incorporación.
  • MAX_ROWS: La cantidad máxima de filas por RPC
  • EMBEDDING_VERSION_COLUMN: Es la columna que administra la versión del modelo de incorporación de textembedding-gecko que se usa para reabastecer tus incorporaciones.
  • MODEL_VERSION: Es la versión del modelo de incorporación de textembedding-gecko.
  • FILTER_CONDITION: Es una condición de filtro particionable que deseas aplicar.

El uso de SAFE.ML.PREDICT muestra NULL para las solicitudes que fallan. También puedes usar SAFE.ML.PREDICT en combinación con un filtro WHERE embedding_column IS NULL para volver a ejecutar la consulta sin calcular las incorporaciones 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/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;

Reemplaza lo siguiente:

  • TABLE_NAME: Es el nombre de la tabla con los datos textuales.
  • EMBEDDING_COLUMN_NAME: El nombre de la columna en la que deseas agregar las incorporaciones generadas
  • DATA_COLUMN_NAME: Es el nombre de la columna con los datos textuales.
  • PROJECT: Es el proyecto que aloja el extremo de Vertex AI.
  • LOCATION: la ubicación del extremo de Vertex AI
  • MODEL_VERSION: Es la versión del modelo de incorporación de textembedding-gecko.
  • MAX_ROWS: La cantidad máxima de filas por RPC
  • EMBEDDING_VERSION_COLUMN: Es la columna que administra la versión del modelo de incorporación de textembedding-gecko que se usa para reabastecer tus incorporaciones.
  • FILTER_CONDITION: Es una condición de filtro particionable que deseas aplicar.

Ejemplo de consulta de reabastecimiento para la tabla 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 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

Ten en cuenta lo siguiente para optimizar el rendimiento cuando reabastezcas 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 la DML particionada. Si las solicitudes de la API de Vertex AI se regulan debido a los límites de cuota de la API de Vertex AI, Spanner vuelve a intentar estas fallas en el modo de transacción de DML particionada por 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 puedes ajustar el paralelismo con la sugerencia a nivel de la sentencia @{pdml_max_parallelism=DESIRED_NUMBER} mientras usas Google SQL. 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 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;

Tamaño del texto en la columna de datos

El modelo de incorporación de Vertex AI tiene límites en la cantidad máxima de tokens para cada entrada de texto. Las diferentes versiones de modelos tienen diferentes límites de tokens. Cada solicitud de Vertex AI puede tener varios campos de texto de entrada, pero existe un límite para la cantidad máxima de tokens presentes 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 que se enviaron a Vertex AI

Puedes usar la sugerencia de consulta @{remote_udf_max_outstanding_rpcs} para aumentar o disminuir la cantidad de solicitudes que se envían a Vertex AI 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 que se envían a Vertex AI desde Spanner con el panel de estadísticas del sistema.

¿Qué sigue?