Generate text embeddings by using the Vertex AI API

Run in Google Colab View source on GitHub

Text embeddings are a way to represent text as numerical vectors. This process lets computers understand and process text data, which is essential for many natural language processing (NLP) tasks.

The following NLP tasks use embeddings:

  • Semantic search: Find documents or passages that are relevant to a query when the query doesn't use the exact same words as the documents.
  • Text classification: Categorize text data into different classes, such as spam and not spam, or positive sentiment and negative sentiment.
  • Machine translation: Translate text from one language to another and preserve the meaning.
  • Text summarization: Create shorter summaries of text.

This notebook uses the Vertex AI text-embeddings API to generate text embeddings that use Google’s large generative artificial intelligence (AI) models. To generate text embeddings by using the Vertex AI text-embeddings API, use MLTransform with the VertexAITextEmbeddings class to specify the model configuration. For more information, see Get text embeddings in the Vertex AI documentation.

For more information about using MLTransform, see Preprocess data with MLTransform in the Apache Beam documentation.

Requirements

To use the Vertex AI text-embeddings API, complete the following prerequisites:

  • Install the google-cloud-aiplatform Python package.
  • Do one of the following tasks:

To use your Google Cloud account, authenticate this notebook.

from google.colab import auth
auth.authenticate_user()

# Replace <PROJECT_ID> with a valid Google Cloud project ID.
project = '<PROJECT_ID>' # @param {type:'string'}

Install dependencies

Install Apache Beam and the dependencies required for the Vertex AI text-embeddings API.

 pip install apache_beam[interactive,gcp]>=2.53.0 --quiet
 
import tempfile
import