feat: add generativeaionvertexai_gemma2_predict_tpu

Author: Joanna Grycz

generative-ai/snippets/gemma2PredictTpu.js

@@ -0,0 +1,75 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// [START generativeaionvertexai_gemma2_predict_tpu]
+// Imports the Google Cloud Prediction Service Client library
+const {
+  // TODO(developer): Uncomment PredictionServiceClient before running the sample.
+  // PredictionServiceClient,
+  helpers,
+} = require('@google-cloud/aiplatform');
+
+async function gemma2PredictTpu(predictionServiceClient) {
+  /**
+   * TODO(developer): Update these variables before running the sample.
+   */
+  const projectId = 'your-project-id';
+  const endpointRegion = 'your-vertex-endpoint-region';
+  const endpointId = 'your-vertex-endpoint-id';
+
+  // Prompt used in the prediction
+  const prompt = 'Why is the sky blue?';
+
+  // Encapsulate the prompt in a correct format for TPUs
+  // Example format: [{prompt: 'Why is the sky blue?', temperature: 0.9}]
+  const input = {
+    prompt,
+    // Parameters for default configuration
+    maxOutputTokens: 1024,
+    temperature: 0.9,
+    topP: 1.0,
+    topK: 1,
+  };
+
+  // Convert input message to a list of GAPIC instances for model input
+  const instances = [helpers.toValue(input)];
+
+  // TODO(developer): Uncomment apiEndpoint and predictionServiceClient before running the sample.
+  // const apiEndpoint = `${endpointRegion}-aiplatform.googleapis.com`;
+
+  // Create a client
+  // predictionServiceClient = new PredictionServiceClient({apiEndpoint});
+
+  // Call the Gemma2 endpoint
+  const gemma2Endpoint = `projects/${projectId}/locations/${endpointRegion}/endpoints/${endpointId}`;
+
+  const [response] = await predictionServiceClient.predict({
+    endpoint: gemma2Endpoint,
+    instances,
+  });
+
+  const predictions = response.predictions;
+  const text = predictions[0].stringValue;
+
+  console.log('Predictions:', text);
+  // [END generativeaionvertexai_gemma2_predict_tpu]
+  return text;
+}
+
+module.exports = gemma2PredictTpu;
+
+gemma2PredictTpu(...process.argv.slice(2)).catch(err => {
+  console.error(err.message);
+  process.exitCode = 1;
+});

generative-ai/snippets/test/gemma2Prediction.test.js

@@ -20,8 +20,9 @@ const {expect} = require('chai');
 const {afterEach, describe, it} = require('mocha');
 const sinon = require('sinon');
 const gemma2PredictGpu = require('../gemma2PredictGpu.js');
+const gemma2PredictTpu = require('../gemma2PredictTpu.js');
 
-const text = `The sky appears blue due to a phenomenon called **Rayleigh scattering**.
+const gpuResponse = `The sky appears blue due to a phenomenon called **Rayleigh scattering**.
 **Here's how it works:**
 1. **Sunlight:** Sunlight is composed of all the colors of the rainbow.
 2. **Earth's Atmosphere:** When sunlight enters the Earth's atmosphere, it collides with tiny particles like nitrogen and oxygen molecules.
@@ -32,26 +33,45 @@ const text = `The sky appears blue due to a phenomenon called **Rayleigh scatter
 * **Longer wavelengths** like red, orange, and yellow are scattered less and travel more directly through the atmosphere. This is why we see these colors during sunrise and sunset, when sunlight has to travel through more of the atmosphere.
 `;
 
-describe('Gemma2 predictions', async () => {
-  const predictionServiceClientMock = {
-    predict: sinon.stub().resolves([
-      {
-        predictions: [
-          {
-            stringValue: text,
-          },
-        ],
-      },
-    ]),
-  };
+const tpuResponse =
+  'The sky appears blue due to a phenomenon called **Rayleigh scattering**.';
 
+describe('Gemma2 predictions', async () => {
   afterEach(() => {
     sinon.restore();
   });
 
   it('should run interference with GPU', async () => {
+    const predictionServiceClientMock = {
+      predict: sinon.stub().resolves([
+        {
+          predictions: [
+            {
+              stringValue: gpuResponse,
+            },
+          ],
+        },
+      ]),
+    };
     const output = await gemma2PredictGpu(predictionServiceClientMock);
 
     expect(output).include('Rayleigh scattering');
   });
+
+  it('should run interference with TPU', async () => {
+    const predictionServiceClientMock = {
+      predict: sinon.stub().resolves([
+        {
+          predictions: [
+            {
+              stringValue: tpuResponse,
+            },
+          ],
+        },
+      ]),
+    };
+    const output = await gemma2PredictTpu(predictionServiceClientMock);
+
+    expect(output).include('Rayleigh scattering');
+  });
 });