Generate text embedding

  from 
  
 __future__ 
  
 import 
 annotations 
 from 
  
 vertexai.language_models 
  
 import 
 TextEmbeddingInput 
 , 
 TextEmbeddingModel 
 def 
  
 embed_text 
 () 
 - 
> list 
 [ 
 list 
 [ 
 float 
 ]]: 
  
 """Embeds texts with a pre-trained, foundational model. 
 Returns: 
 A list of lists containing the embedding vectors for each input text 
 """ 
 # A list of texts to be embedded. 
 texts 
 = 
 [ 
 "banana muffins? " 
 , 
 "banana bread? banana muffins?" 
 ] 
 # The dimensionality of the output embeddings. 
 dimensionality 
 = 
 3072 
 # The task type for embedding. Check the available tasks in the model's documentation. 
 task 
 = 
 "RETRIEVAL_DOCUMENT" 
 model 
 = 
 TextEmbeddingModel 
 . 
 from_pretrained 
 ( 
 "gemini-embedding-001" 
 ) 
 kwargs 
 = 
 dict 
 ( 
 output_dimensionality 
 = 
 dimensionality 
 ) 
 if 
 dimensionality 
 else 
 {} 
 embeddings 
 = 
 [] 
 # gemini-embedding-001 takes one input at a time 
 for 
 text 
 in 
 texts 
 : 
 text_input 
 = 
 TextEmbeddingInput 
 ( 
 text 
 , 
 task 
 ) 
 embedding 
 = 
 model 
 . 
 get_embeddings 
 ([ 
 text_input 
 ], 
 ** 
 kwargs 
 ) 
 print 
 ( 
 embedding 
 ) 
 # Example response: 
 # [[0.006135190837085247, -0.01462465338408947, 0.004978656303137541, ...]] 
 embeddings 
 . 
 append 
 ( 
 embedding 
 [ 
 0 
 ] 
 . 
 values 
 ) 
 return 
 embeddings