Generate Embeddings from text using Batch processing

The code sample showcases how to use a pre-trained model to batch generate embeddings for a list of text inputs, and store them in a specified location.

Code sample

Java

Before trying this sample, follow the Java setup instructions in the Vertex AI quickstart using client libraries . For more information, see the Vertex AI Java API reference documentation .

To authenticate to Vertex AI, set up Application Default Credentials. For more information, see Set up authentication for a local development environment .

  import 
  
 com.google.cloud.aiplatform.v1. BatchPredictionJob 
 
 ; 
 import 
  
 com.google.cloud.aiplatform.v1. GcsDestination 
 
 ; 
 import 
  
 com.google.cloud.aiplatform.v1. GcsSource 
 
 ; 
 import 
  
 com.google.cloud.aiplatform.v1. JobServiceClient 
 
 ; 
 import 
  
 com.google.cloud.aiplatform.v1. JobServiceSettings 
 
 ; 
 import 
  
 com.google.cloud.aiplatform.v1. LocationName 
 
 ; 
 import 
  
 java.io.IOException 
 ; 
 public 
  
 class 
 EmbeddingBatchSample 
  
 { 
  
 public 
  
 static 
  
 void 
  
 main 
 ( 
 String 
 [] 
  
 args 
 ) 
  
 throws 
  
 IOException 
 , 
  
 InterruptedException 
  
 { 
  
 // TODO(developer): Replace these variables before running the sample. 
  
 String 
  
 project 
  
 = 
  
 "YOUR_PROJECT_ID" 
 ; 
  
 String 
  
 location 
  
 = 
  
 "us-central1" 
 ; 
  
 // inputUri: URI of the input dataset. 
  
 // Could be a BigQuery table or a Google Cloud Storage file. 
  
 // E.g. "gs://[BUCKET]/[DATASET].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]" 
  
 String 
  
 inputUri 
  
 = 
  
 "gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl" 
 ; 
  
 // outputUri: URI where the output will be stored. 
  
 // Could be a BigQuery table or a Google Cloud Storage file. 
  
 // E.g. "gs://[BUCKET]/[OUTPUT].jsonl" OR "bq://[PROJECT].[DATASET].[TABLE]" 
  
 String 
  
 outputUri 
  
 = 
  
 "gs://YOUR_BUCKET/embedding_batch_output" 
 ; 
  
 String 
  
 textEmbeddingModel 
  
 = 
  
 "text-embedding-005" 
 ; 
  
 embeddingBatchSample 
 ( 
 project 
 , 
  
 location 
 , 
  
 inputUri 
 , 
  
 outputUri 
 , 
  
 textEmbeddingModel 
 ); 
  
 } 
  
 // Generates embeddings from text using batch processing. 
  
 // Read more: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings 
  
 public 
  
 static 
  
  BatchPredictionJob 
 
  
 embeddingBatchSample 
 ( 
  
 String 
  
 project 
 , 
  
 String 
  
 location 
 , 
  
 String 
  
 inputUri 
 , 
  
 String 
  
 outputUri 
 , 
  
 String 
  
 textEmbeddingModel 
 ) 
  
 throws 
  
 IOException 
  
 { 
  
  BatchPredictionJob 
 
  
 response 
 ; 
  
  JobServiceSettings 
 
  
 jobServiceSettings 
  
 = 
  
  JobServiceSettings 
 
 . 
 newBuilder 
 () 
  
 . 
 setEndpoint 
 ( 
 "us-central1-aiplatform.googleapis.com:443" 
 ). 
 build 
 (); 
  
  LocationName 
 
  
 parent 
  
 = 
  
  LocationName 
 
 . 
 of 
 ( 
 project 
 , 
  
 location 
 ); 
  
 String 
  
 modelName 
  
 = 
  
 String 
 . 
 format 
 ( 
 "projects/%s/locations/%s/publishers/google/models/%s" 
 , 
  
 project 
 , 
  
 location 
 , 
  
 textEmbeddingModel 
 ); 
  
 // Initialize client that will be used to send requests. This client only needs to be created 
  
 // once, and can be reused for multiple requests. 
  
 try 
  
 ( 
  JobServiceClient 
 
  
 client 
  
 = 
  
  JobServiceClient 
 
 . 
 create 
 ( 
 jobServiceSettings 
 )) 
  
 { 
  
  BatchPredictionJob 
 
  
 batchPredictionJob 
  
 = 
  
  BatchPredictionJob 
 
 . 
 newBuilder 
 () 
  
 . 
 setDisplayName 
 ( 
 "my embedding batch job " 
  
 + 
  
 System 
 . 
 currentTimeMillis 
 ()) 
  
 . 
 setModel 
 ( 
 modelName 
 ) 
  
 . 
  setInputConfig 
 
 ( 
  
  BatchPredictionJob 
 
 . 
 InputConfig 
 . 
 newBuilder 
 () 
  
 . 
 setGcsSource 
 ( 
  GcsSource 
 
 . 
 newBuilder 
 (). 
  addUris 
 
 ( 
 inputUri 
 ). 
 build 
 ()) 
  
 . 
  setInstancesFormat 
 
 ( 
 "jsonl" 
 ) 
  
 . 
 build 
 ()) 
  
 . 
 setOutputConfig 
 ( 
  
  BatchPredictionJob 
 
 . 
 OutputConfig 
 . 
 newBuilder 
 () 
  
 . 
 setGcsDestination 
 ( 
  GcsDestination 
 
 . 
 newBuilder 
 () 
  
 . 
  setOutputUriPrefix 
 
 ( 
 outputUri 
 ). 
 build 
 ()) 
  
 . 
  setPredictionsFormat 
 
 ( 
 "jsonl" 
 ) 
  
 . 
 build 
 ()) 
  
 . 
 build 
 (); 
  
 response 
  
 = 
  
 client 
 . 
 createBatchPredictionJob 
 ( 
 parent 
 , 
  
 batchPredictionJob 
 ); 
  
 System 
 . 
 out 
 . 
 format 
 ( 
 "response: %s\n" 
 , 
  
 response 
 ); 
  
 System 
 . 
 out 
 . 
 format 
 ( 
 "\tName: %s\n" 
 , 
  
 response 
 . 
  getName 
 
 ()); 
  
 } 
  
 return 
  
 response 
 ; 
  
 } 
 } 
 

Python

Before trying this sample, follow the Python setup instructions in the Vertex AI quickstart using client libraries . For more information, see the Vertex AI Python API reference documentation .

To authenticate to Vertex AI, set up Application Default Credentials. For more information, see Set up authentication for a local development environment .

  import 
  
  vertexai 
 
 from 
  
 vertexai.preview 
  
 import 
  language_models 
 
 # TODO(developer): Update & uncomment line below 
 # PROJECT_ID = "your-project-id" 
  vertexai 
 
 . 
 init 
 ( 
 project 
 = 
 PROJECT_ID 
 , 
 location 
 = 
 "us-central1" 
 ) 
 input_uri 
 = 
 ( 
 "gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl" 
 ) 
 # Format: `"gs://your-bucket-unique-name/directory/` or `bq://project_name.llm_dataset` 
 output_uri 
 = 
 OUTPUT_URI 
 textembedding_model 
 = 
  language_models 
 
 . 
 TextEmbeddingModel 
 . 
 from_pretrained 
 ( 
 "textembedding-gecko@003" 
 ) 
 batch_prediction_job 
 = 
 textembedding_model 
 . 
 batch_predict 
 ( 
 dataset 
 = 
 [ 
 input_uri 
 ], 
 destination_uri_prefix 
 = 
 output_uri 
 , 
 ) 
 print 
 ( 
 batch_prediction_job 
 . 
 display_name 
 ) 
 print 
 ( 
 batch_prediction_job 
 . 
 resource_name 
 ) 
 print 
 ( 
 batch_prediction_job 
 . 
 state 
 ) 
 # Example response: 
 # BatchPredictionJob 2024-09-10 15:47:51.336391 
 # projects/1234567890/locations/us-central1/batchPredictionJobs/123456789012345 
 # JobState.JOB_STATE_SUCCEEDED 
 

What's next

To search and filter code samples for other Google Cloud products, see the Google Cloud sample browser .

Design a Mobile Site
View Site in Mobile | Classic
Share by: