Text embeddings API

from__future__import annotations
fromvertexai.language_modelsimport TextEmbeddingInput, TextEmbeddingModel
defembed_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

import(
"context"
"fmt"
"io"
aiplatform"cloud.google.com/go/aiplatform/apiv1"
"cloud.google.com/go/aiplatform/apiv1/aiplatformpb"
"google.golang.org/api/option"
"google.golang.org/protobuf/types/known/structpb"
)
// embedTexts shows how embeddings are set for gemini-embedding-001 model
funcembedTexts(wio.Writer,project,locationstring)error{
// location := "us-central1"
ctx:=context.Background()
apiEndpoint:=fmt.Sprintf("%s-aiplatform.googleapis.com:443",location)
dimensionality:=3072
model:="gemini-embedding-001"
texts:=[]string{"banana muffins? ","banana bread? banana muffins?"}
client,err:=aiplatform.NewPredictionClient (ctx,option.WithEndpoint(apiEndpoint))
iferr!=nil{
returnerr
}
deferclient.Close()
endpoint:=fmt.Sprintf("projects/%s/locations/%s/publishers/google/models/%s",project,location,model)
allEmbeddings:=make([][]float32,0,len(texts))
// gemini-embedding-001 takes 1 input at a time
for_,text:=rangetexts{
instances:=make([]*structpb.Value,1)
instances[0]=structpb.NewStructValue(&structpb.Struct{
Fields:map[string]*structpb.Value{
"content":structpb.NewStringValue(text),
"task_type":structpb.NewStringValue("QUESTION_ANSWERING"),
},
})
params:=structpb.NewStructValue(&structpb.Struct{
Fields:map[string]*structpb.Value{
"outputDimensionality":structpb.NewNumberValue(float64(dimensionality)),
},
})
req:=&aiplatformpb.PredictRequest{
Endpoint:endpoint,
Instances:instances,
Parameters:params,
}
resp,err:=client.Predict(ctx,req)
iferr!=nil{
returnerr
}
// Process the prediction for the single text
// The response will contain one prediction because we sent one instance.
iflen(resp.Predictions)==0{
returnfmt.Errorf("no predictions returned for text \"%s\"",text)
}
prediction:=resp.Predictions[0]
embeddingValues:=prediction.GetStructValue().Fields["embeddings"].GetStructValue().Fields["values"].GetListValue().Values
currentEmbedding:=make([]float32,len(embeddingValues))
forj,value:=rangeembeddingValues{
currentEmbedding[j]=float32(value.GetNumberValue())
}
allEmbeddings=append(allEmbeddings,currentEmbedding)
}
iflen(allEmbeddings) > 0{
fmt.Fprintf(w,"Dimensionality: %d. Embeddings length: %d",len(allEmbeddings[0]),len(allEmbeddings))
}else{
fmt.Fprintln(w,"No texts were processed.")
}
returnnil
}

import staticjava.util.stream.Collectors.toList;
importcom.google.cloud.aiplatform.v1.EndpointName ;
importcom.google.cloud.aiplatform.v1.PredictRequest ;
importcom.google.cloud.aiplatform.v1.PredictResponse ;
importcom.google.cloud.aiplatform.v1.PredictionServiceClient ;
importcom.google.cloud.aiplatform.v1.PredictionServiceSettings ;
importcom.google.protobuf.Struct ;
importcom.google.protobuf.Value ;
importjava.io.IOException;
importjava.util.ArrayList;
importjava.util.List;
importjava.util.OptionalInt;
importjava.util.regex.Matcher;
importjava.util.regex.Pattern;
publicclass PredictTextEmbeddingsSample{
publicstaticvoidmain(String[]args)throwsIOException{
// TODO(developer): Replace these variables before running the sample.
// Details about text embedding request structure and supported models are available in:
// https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings
Stringendpoint="us-central1-aiplatform.googleapis.com:443";
Stringproject="YOUR_PROJECT_ID";
Stringmodel="gemini-embedding-001";
predictTextEmbeddings(
endpoint,
project,
model,
List.of("banana bread?","banana muffins?"),
"QUESTION_ANSWERING",
OptionalInt.of(3072));
}
// Gets text embeddings from a pretrained, foundational model.
publicstaticList<List<Float>>predictTextEmbeddings(
Stringendpoint,
Stringproject,
Stringmodel,
List<String>texts,
Stringtask,
OptionalIntoutputDimensionality)
throwsIOException{
PredictionServiceSettings settings=
PredictionServiceSettings .newBuilder().setEndpoint(endpoint).build();
Matchermatcher=Pattern.compile("^(?<Location>\\w+-\\w+)").matcher(endpoint);
Stringlocation=matcher.matches()?matcher.group("Location"):"us-central1";
EndpointName endpointName=
EndpointName .ofProjectLocationPublisherModelName (project,location,"google",model);
List<List<Float>>floats=newArrayList<>();
// You can use this prediction service client for multiple requests.
try(PredictionServiceClient client=PredictionServiceClient .create(settings)){
// gemini-embedding-001 takes one input at a time.
for(inti=0;i < texts.size();i++){
PredictRequest .Builderrequest=
PredictRequest .newBuilder().setEndpoint(endpointName.toString ());
if(outputDimensionality.isPresent()){
request.setParameters(
Value .newBuilder()
.setStructValue(
Struct .newBuilder()
.putFields(
"outputDimensionality",valueOf(outputDimensionality.getAsInt()))
.build()));
}
request.addInstances(
Value .newBuilder()
.setStructValue(
Struct .newBuilder()
.putFields("content",valueOf(texts.get(i)))
.putFields("task_type",valueOf(task))
.build()));
PredictResponse response=client.predict(request.build());
for(Value prediction:response.getPredictionsList ()){
Value embeddings=prediction.getStructValue().getFieldsOrThrow("embeddings");
Value values=embeddings.getStructValue ().getFieldsOrThrow("values");
floats.add(
values.getListValue ().getValuesList().stream()
.map(Value::getNumberValue)
.map(Double::floatValue)
.collect(toList()));
}
}
returnfloats;
}
}
privatestaticValue valueOf(Strings){
returnValue .newBuilder().setStringValue(s).build();
}
privatestaticValue valueOf(intn){
returnValue .newBuilder().setNumberValue(n).build();
}
}

asyncfunctionmain(
project,
model='gemini-embedding-001',
texts='banana bread?;banana muffins?',
task='QUESTION_ANSWERING',
dimensionality=0,
apiEndpoint='us-central1-aiplatform.googleapis.com'
){
constaiplatform=require('@google-cloud/aiplatform');
const{PredictionServiceClient}=aiplatform.v1;
const{helpers}=aiplatform;// helps construct protobuf.Value objects.
constclientOptions={apiEndpoint:apiEndpoint};
constlocation='us-central1';
constendpoint=`projects/${project}/locations/${location}/publishers/google/models/${model}`;
asyncfunctioncallPredict(){
constinstances=texts
.split(';')
.map(e=>helpers .toValue({content:e,task_type:task}));
constclient=newPredictionServiceClient (clientOptions);
constparameters=helpers .toValue(
dimensionality > 0?{outputDimensionality:parseInt(dimensionality)}:{}
);
constallEmbeddings=[]
// gemini-embedding-001 takes one input at a time.
for(constinstanceofinstances){
constrequest={endpoint,instances:[instance],parameters};
const[response]=awaitclient.predict(request);
constpredictions=response.predictions;
constembeddings=predictions.map(p=>{
constembeddingsProto=p.structValue.fields.embeddings;
constvaluesProto=embeddingsProto.structValue.fields.values;
returnvaluesProto.listValue.values.map(v=>v.numberValue);
});
allEmbeddings.push(embeddings[0])
}
console.log('Got embeddings: \n'+JSON.stringify(allEmbeddings));
}
callPredict();
}