feat: Elastic Text-Embedding Model demo. (override 003 to 004)

GoogleCloudPlatform · May 8, 2024 · 20f37c5 · 20f37c5
1 parent b45f484
commit 20f37c5
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 7 deletions.
diff --git a/aiplatform/snippets/embeddings.go b/aiplatform/snippets/embeddings.go
@@ -28,7 +28,8 @@ import (
 )
 
 func embedTexts(
-	apiEndpoint, project, model string, texts []string, task string) ([][]float32, error) {
+	apiEndpoint, project, model string, texts []string,
+	task string, customOutputDimensionality *int) ([][]float32, error) {
 	ctx := context.Background()
 
 	client, err := aiplatform.NewPredictionClient(ctx, option.WithEndpoint(apiEndpoint))
@@ -52,10 +53,18 @@ func embedTexts(
 			},
 		})
 	}
+	outputDimensionality := structpb.NewNullValue()
+	if customOutputDimensionality != nil {
+		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 {

diff --git a/aiplatform/snippets/embeddings_test.go b/aiplatform/snippets/embeddings_test.go
@@ -23,14 +23,15 @@ import (
 func TestGenerateEmbeddings(t *testing.T) {
 	tc := testutil.SystemTest(t)
 	apiEndpoint := "us-central1-aiplatform.googleapis.com:443"
-	model := "textembedding-gecko@003"
+	model := "text-embedding-004"
 	texts := []string{"banana muffins? ", "banana bread? banana muffins?"}
-	embeddings, err := embedTexts(apiEndpoint, tc.ProjectID, model, texts, "RETRIEVAL_DOCUMENT")
+	dimensionality := 5
+	embeddings, err := embedTextsPreview(apiEndpoint, tc.ProjectID, model, texts, "QUESTION_ANSWERING", &dimensionality)
 	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))
+	if len(embeddings) != len(texts) || len(embeddings[0]) != dimensionality {
+		t.Errorf("len(embeddings), len(embeddings[0]) = %d, %d, want %d, %d", len(embeddings), len(embeddings[0]), len(texts), dimensionality)
 	}
 }