feat(vertexai): Elastic Text-Embedding Model demo.

aiplatform/snippets/embeddings.go

@@ -28,17 +28,21 @@ import (
 )
 
 func embedTexts(
-	apiEndpoint, project, model string, texts []string, task string) ([][]float32, error) {
+	project, location string, texts []string) ([][]float32, error) {
 	ctx := context.Background()
 
+	apiEndpoint := fmt.Sprintf("%s-aiplatform.googleapis.com:443", location)
+	model := "text-embedding-004"
+	task := "QUESTION_ANSWERING"
+	customOutputDimensionality := 5
+
 	client, err := aiplatform.NewPredictionClient(ctx, option.WithEndpoint(apiEndpoint))
 	if err != nil {
 		return nil, err
 	}
 	defer client.Close()
 
 	match := regexp.MustCompile(`^(\w+-\w+)`).FindStringSubmatch(apiEndpoint)
-	location := "us-central1"
 	if match != nil {
 		location = match[1]
 	}
@@ -52,10 +56,17 @@ func embedTexts(
 			},
 		})
 	}
+	outputDimensionality := structpb.NewNullValue()
+	outputDimensionality = structpb.NewNumberValue(float64(customOutputDimensionality))
+
+	params := structpb.NewStructValue(&structpb.Struct{
+		Fields: map[string]*structpb.Value{"outputDimensionality": outputDimensionality},
+	})
 
 	req := &aiplatformpb.PredictRequest{
-		Endpoint:  endpoint,
-		Instances: instances,
+		Endpoint:   endpoint,
+		Instances:  instances,
+		Parameters: params,
 	}
 	resp, err := client.Predict(ctx, req)
 	if err != nil {

aiplatform/snippets/embeddings_test.go

@@ -22,15 +22,20 @@ import (
 
 func TestGenerateEmbeddings(t *testing.T) {
 	tc := testutil.SystemTest(t)
-	apiEndpoint := "us-central1-aiplatform.googleapis.com:443"
-	model := "textembedding-gecko@003"
 	texts := []string{"banana muffins? ", "banana bread? banana muffins?"}
-	embeddings, err := embedTexts(apiEndpoint, tc.ProjectID, model, texts, "RETRIEVAL_DOCUMENT")
+	dimensionality := 5
+	location := "us-central1"
+	embeddings, err := embedTexts(tc.ProjectID, location, texts)
 	if err != nil {
 		t.Fatal(err)
 	}
-	if len(embeddings) != len(texts) || len(embeddings[0]) != 768 {
-		t.Errorf("len(embeddings), len(embeddings[0]) = %d, %d, want %d, 768", len(embeddings), len(embeddings[0]), len(texts))
+
+	embeddingsLen := len(embeddings)
+	textsLen := len(texts)
+	embeddingDimensionality := len(embeddings[0])
+
+	if embeddingsLen != textsLen || embeddingDimensionality != dimensionality {
+		t.Errorf("embeddingsLen, embeddingDimensionality = %d, %d, want %d, %d", embeddingsLen, embeddingDimensionality, textsLen, dimensionality)
 	}
 }