Get Vertex AI text embeddings

A text embedding is a vector representation of text data, and they are used in many ways to find similar items. You interact with them every time you complete a Google search or see recommendations when shopping online. When you create text embeddings, you get vector representations of natural text as arrays of floating point numbers. This means that all of your input text is assigned a numerical representation. By comparing the numerical distance between the vector representations of two pieces of text, an application can determine the similarity between the text or the objects represented by the text.

With the Vertex AI text-embeddings API, you can create a text embedding with Generative AI. Using this tutorial you can generate text-embeddings for the data stored in Spanner and Vertex AI embedding models like the textembedding-gecko model.

To learn more about embeddings, see Get text embeddings.

Objective

In this tutorial, you learn how to:

  • Register the Vertex AI textembedding-gecko model in a Spanner schema using DDL statements.
  • Reference the registered model using SQL queries to generate embeddings from data stored in Spanner.

Costs

This tutorial uses billable components of Google Cloud, including:

  • Spanner
  • Vertex AI

For more information about Spanner costs, see the Spanner pricing page.

For more information about Vertex AI costs, see the Vertex AI pricing page.

Generate and store text embeddings

Depending on the model you use, generating embeddings might take some time. For more performance sensitive workloads, the best practice is to avoid generating embeddings in read-write transactions. Instead, generate the embeddings in a read-only transaction using the following SQL examples.

GoogleSQL

Register a text embeddings model in Spanner

In GoogleSQL, you must register a model before using it with the ML.PREDICT function. To register the textembedding-gecko model in a Spanner database, execute the following DDL statement:

CREATE MODEL MODEL_NAME
INPUT(content STRING(MAX))
OUTPUT(
  embeddings
    STRUCT<
      statistics STRUCT<truncated BOOL, token_count FLOAT64>,
      values ARRAY<FLOAT64>>
)
REMOTE OPTIONS (
  endpoint = '//aiplatform.googleapis.com/projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko$MODEL_VERSION'
);

Replace the following:

  • MODEL_NAME: the name of the embedding model
  • PROJECT: the project hosting the Vertex AI endpoint
  • LOCATION: the location of the Vertex AI endpoint
  • MODEL_VERSION: the version of the textembedding-gecko embedding model

Spanner grants appropriate permissions automatically. If it doesn't, review the model endpoint access control.

Schema discovery and validation is not available for Generative AI models. You are required to provide INPUT and OUTPUT clauses which match against the models schema. For the full schema of the Gecko model, see Get text embeddings.

Generate text embeddings

To generate embeddings, pass a piece of text directly to the ML.PREDICT function using the following SQL:

SELECT embeddings.values
FROM ML.PREDICT(
  MODEL MODEL_NAME,
  (SELECT "A product description" as content)
);

To generate embeddings for data stored in a table, use the following SQL:

SELECT id, embeddings.values
FROM ML.PREDICT(
  MODEL MODEL_NAME,
  (SELECT id, description as content FROM Products)
);

Store text embeddings

After generating the embeddings in a read-only transaction, store them in Spanner so they can be managed with your operational data. To store the embeddings, use a read-write transaction.

For workloads that are less performance sensitive, you can generate and insert embeddings with the following SQL in a read-write transaction:

CREATE TABLE Products(
  id INT64 NOT NULL,
  description STRING(MAX),
  embeddings ARRAY<FLOAT32>,
) PRIMARY KEY(id);

INSERT INTO Products (id, description, embeddings)
SELECT @Id, @Description, embeddings.values
FROM ML.PREDICT(
  MODEL MODEL_NAME,
  (SELECT @Description as content)
);

PostgreSQL

Generate text embeddings

To generate embeddings, pass a piece of text directly to the spanner.ML_PREDICT_ROW function using the following SQL:

SELECT
  spanner.ML_PREDICT_ROW(
    'projects/PROJECT/locations/LOCATION/publishers/google/models/textembedding-gecko$MODEL_VERSION',
    '{"instances": [{"content": "A product description"}]}'::jsonb
  ) ->'predictions'->0->'embeddings'->'values';

Replace the following:

  • PROJECT: the project hosting the Vertex AI endpoint
  • LOCATION: the location of the Vertex AI endpoint
  • MODEL_VERSION: the version of the textembedding-gecko embedding model

To generate embeddings for data stored in a table, use the following SQL:

SELECT id, 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))))
  ) -> `predictions`->0->`embeddings`->`values``
FROM Products;

Replace the following:

  • PROJECT: the project hosting the Vertex AI endpoint
  • LOCATION: the location of the Vertex AI endpoint
  • MODEL_VERSION: the version of the textembedding-gecko embedding model

Store text embeddings

After generating the embeddings in a read-only transaction, store them in Spanner so they can be managed with your operational data. To store the embeddings, use a read-write transaction.

For workloads that are less performance sensitive, you can generate and insert embeddings with the following SQL in a read-write transaction:

CREATE TABLE Products (
  id INT8 NOT NULL,
  description TEXT,
  embeddings REAL[],
  PRIMARY KEY(id)
);

INSERT INTO Products (id, description, embeddings)
SELECT @Id, @Description, 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', @Description)))
  ) -> 'predictions'->0->'embeddings'->'values'
));

Replace the following:

  • PROJECT: the project hosting the Vertex AI endpoint
  • LOCATION: the location of the Vertex AI endpoint
  • MODEL_VERSION: the version of the textembedding-gecko embedding model

Update text embeddings

To update your embeddings or to ingest data in realtime, use the UPDATE (GoogleSQL and PostgreSQL) statement.

To update the Products table in the previous example, use the following SQL:

GoogleSQL 

UPDATE Products
SET
  description = @description,
  embeddings = (SELECT embeddings.values
                  FROM ML.PREDICT(MODEL MODEL_NAME, (SELECT @description as content))
              )
WHERE id = @id;

Replace the following:

  • MODEL_NAME: the name of the embedding model

PostgreSQL 

UPDATE
  Products
SET
  description = $1,
  embeddings = 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', $1)))
    ) -> 'predictions'->0->'embeddings'->'values')
WHERE
  id = $2;

Replace the following:

  • PROJECT: the project hosting the Vertex AI endpoint
  • LOCATION: the location of the Vertex AI endpoint
  • MODEL_VERSION: the version of the textembedding-gecko embedding model

What's next