correcting according to comments

Author: BigBlackWolf

gemma2/gemma2_predict_gpu.go

@@ -26,11 +26,6 @@ import (
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
-type ClientInterface interface {
-	Close() error
-	Predict(ctx context.Context, req *aiplatformpb.PredictRequest, opts ...gax.CallOption) (*aiplatformpb.PredictResponse, error)
-}
-
 // predictGPU demonstrates how to run interference on a Gemma2 model deployed to a Vertex AI endpoint with GPU accelerators.
 func predictGPU(w io.Writer, client ClientInterface, projectID, location, endpointID string) error {
 	ctx := context.Background()
@@ -54,6 +49,7 @@ func predictGPU(w io.Writer, client ClientInterface, projectID, location, endpoi
 
 	// Encapsulate the prompt in a correct format for TPUs.
 	// Pay attention that prompt should be set in "inputs" field.
+	// Example format: [{'inputs': 'Why is the sky blue?', 'parameters': {'temperature': 0.9}}]
 	promptValue, err := structpb.NewValue(map[string]interface{}{
 		"inputs":     prompt,
 		"parameters": parameters,
@@ -81,3 +77,8 @@ func predictGPU(w io.Writer, client ClientInterface, projectID, location, endpoi
 }
 
 // [END generativeaionvertexai_gemma2_predict_gpu]
+
+type ClientInterface interface {
+	Close() error
+	Predict(ctx context.Context, req *aiplatformpb.PredictRequest, opts ...gax.CallOption) (*aiplatformpb.PredictResponse, error)
+}

gemma2/gemma2_predict_tpu.go

@@ -26,11 +26,6 @@ import (
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
-type PredictClientInterface interface {
-	Close() error
-	Predict(ctx context.Context, req *aiplatformpb.PredictRequest, opts ...gax.CallOption) (*aiplatformpb.PredictResponse, error)
-}
-
 // predictTPU demonstrates how to run interference on a Gemma2 model deployed to a Vertex AI endpoint with TPU accelerators.
 func predictTPU(w io.Writer, client PredictClientInterface, projectID, location, endpointID string) error {
 	ctx := context.Background()
@@ -53,6 +48,7 @@ func predictTPU(w io.Writer, client PredictClientInterface, projectID, location,
 	}
 
 	// Encapsulate the prompt in a correct format for TPUs.
+	// Example format: [{'prompt': 'Why is the sky blue?', 'temperature': 0.9}]
 	promptValue, err := structpb.NewValue(map[string]interface{}{
 		"prompt":     prompt,
 		"parameters": parameters,
@@ -80,3 +76,8 @@ func predictTPU(w io.Writer, client PredictClientInterface, projectID, location,
 }
 
 // [END generativeaionvertexai_gemma2_predict_tpu]
+
+type PredictClientInterface interface {
+	Close() error
+	Predict(ctx context.Context, req *aiplatformpb.PredictRequest, opts ...gax.CallOption) (*aiplatformpb.PredictResponse, error)
+}

gemma2/gemma2_test.go

@@ -22,7 +22,14 @@ import (
 	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
 )
 
-func TestPredictGPU(t *testing.T) {
+const (
+	GPUEndpointID     string = "123456789"
+	GPUEndpointRegion        = "us-east1"
+	TPUEndpointID            = "987654321"
+	TPUEndpointRegion        = "us-west1"
+)
+
+func TestPredictGemma2(t *testing.T) {
 	tc := testutil.SystemTest(t)
 
 	projectID := tc.ProjectID
@@ -32,9 +39,7 @@ func TestPredictGPU(t *testing.T) {
 	t.Run("GPU predict", func(t *testing.T) {
 		buf.Reset()
 		// Mock ID used to check if GPU was called
-		endpointID := "123456789"
-		location := "us-east4"
-		if err := predictGPU(&buf, client, projectID, location, endpointID); err != nil {
+		if err := predictGPU(&buf, client, projectID, GPUEndpointRegion, GPUEndpointID); err != nil {
 			t.Fatal(err)
 		}
 
@@ -46,9 +51,7 @@ func TestPredictGPU(t *testing.T) {
 	t.Run("TPU predict", func(t *testing.T) {
 		buf.Reset()
 		// Mock ID used to check if TPU was called
-		endpointID := "123456789"
-		location := "us-west1"
-		if err := predictTPU(&buf, client, projectID, location, endpointID); err != nil {
+		if err := predictTPU(&buf, client, projectID, TPUEndpointRegion, TPUEndpointID); err != nil {
 			t.Fatal(err)
 		}
 

gemma2/mock_client.go

@@ -16,6 +16,8 @@ package snippets
 
 import (
 	"context"
+	"fmt"
+	"strings"
 
 	"cloud.google.com/go/aiplatform/apiv1/aiplatformpb"
 	gax "github.com/googleapis/gax-go/v2"
@@ -43,5 +45,31 @@ func (client PredictionsClient) Predict(ctx context.Context, req *aiplatformpb.P
 	response := &aiplatformpb.PredictResponse{
 		Predictions: []*structpb.Value{structpb.NewStringValue(mockedResponse)},
 	}
+
+	instance := req.Instances[0].GetStructValue()
+	if ok := strings.Contains(req.Endpoint, fmt.Sprintf("locations/%s/endpoints/%s", GPUEndpointRegion, GPUEndpointID)); ok {
+		if err := instance.Fields["inputs"].GetStringValue(); err == "" {
+			return nil, fmt.Errorf("invalid request")
+		}
+	} else if ok := strings.Contains(req.Endpoint, fmt.Sprintf("locations/%s/endpoints/%s", TPUEndpointRegion, TPUEndpointID)); ok {
+		if err := instance.Fields["prompt"].GetStringValue(); err == "" {
+			return nil, fmt.Errorf("invalid request")
+		}
+	}
+
+	params := req.Instances[0].GetStructValue().Fields["parameters"].GetStructValue()
+
+	if err := params.Fields["temperature"].GetNumberValue(); err == 0 {
+		return nil, fmt.Errorf("invalid request")
+	}
+	if err := params.Fields["maxOutputTokens"].GetNumberValue(); err == 0 {
+		return nil, fmt.Errorf("invalid request")
+	}
+	if err := params.Fields["topP"].GetNumberValue(); err == 0 {
+		return nil, fmt.Errorf("invalid request")
+	}
+	if err := params.Fields["topK"].GetNumberValue(); err == 0 {
+		return nil, fmt.Errorf("invalid request")
+	}
 	return response, nil
 }