Skip to content

Commit beb84f3

Browse files
gryczjgryczj
authored
feat: gemma2 models samples (#3869)
* feat: gemma2 models samples * feat: add generativeaionvertexai_gemma2_predict_tpu * Add check for request structure
1 parent 7949c97 commit beb84f3

File tree

3 files changed

+287
-0
lines changed

3 files changed

+287
-0
lines changed

‎ai-platform/snippets/gemma2PredictGpu.js

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
// Copyright 2024 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
'use strict';
16+
17+
async function gemma2PredictGpu(predictionServiceClient) {
18+
// [START generativeaionvertexai_gemma2_predict_gpu]
19+
// Imports the Google Cloud Prediction Service Client library
20+
const {
21+
// TODO(developer): Uncomment PredictionServiceClient before running the sample.
22+
// PredictionServiceClient,
23+
helpers,
24+
} = require('@google-cloud/aiplatform');
25+
/**
26+
* TODO(developer): Update these variables before running the sample.
27+
*/
28+
const projectId = 'your-project-id';
29+
const endpointRegion = 'your-vertex-endpoint-region';
30+
const endpointId = 'your-vertex-endpoint-id';
31+
32+
// Default configuration
33+
const config = {maxOutputTokens: 1024, temperature: 0.9, topP: 1.0, topK: 1};
34+
// Prompt used in the prediction
35+
const prompt = 'Why is the sky blue?';
36+
37+
// Encapsulate the prompt in a correct format for GPUs
38+
// Example format: [{inputs: 'Why is the sky blue?', parameters: {temperature: 0.9}}]
39+
const input = {
40+
inputs: prompt,
41+
parameters: config,
42+
};
43+
44+
// Convert input message to a list of GAPIC instances for model input
45+
const instances = [helpers.toValue(input)];
46+
47+
// TODO(developer): Uncomment apiEndpoint and predictionServiceClient before running the sample.
48+
// const apiEndpoint = `${endpointRegion}-aiplatform.googleapis.com`;
49+
50+
// Create a client
51+
// predictionServiceClient = new PredictionServiceClient({apiEndpoint});
52+
53+
// Call the Gemma2 endpoint
54+
const gemma2Endpoint = `projects/${projectId}/locations/${endpointRegion}/endpoints/${endpointId}`;
55+
56+
const [response] = await predictionServiceClient.predict({
57+
endpoint: gemma2Endpoint,
58+
instances,
59+
});
60+
61+
const predictions = response.predictions;
62+
const text = predictions[0].stringValue;
63+
64+
console.log('Predictions:', text);
65+
// [END generativeaionvertexai_gemma2_predict_gpu]
66+
return text;
67+
}
68+
69+
module.exports = gemma2PredictGpu;
70+
71+
// TODO(developer): Uncomment below lines before running the sample.
72+
// gemma2PredictGpu(...process.argv.slice(2)).catch(err => {
73+
// console.error(err.message);
74+
// process.exitCode = 1;
75+
// });

‎ai-platform/snippets/gemma2PredictTpu.js

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
// Copyright 2024 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
'use strict';
16+
17+
async function gemma2PredictTpu(predictionServiceClient) {
18+
// [START generativeaionvertexai_gemma2_predict_tpu]
19+
// Imports the Google Cloud Prediction Service Client library
20+
const {
21+
// TODO(developer): Uncomment PredictionServiceClient before running the sample.
22+
// PredictionServiceClient,
23+
helpers,
24+
} = require('@google-cloud/aiplatform');
25+
/**
26+
* TODO(developer): Update these variables before running the sample.
27+
*/
28+
const projectId = 'your-project-id';
29+
const endpointRegion = 'your-vertex-endpoint-region';
30+
const endpointId = 'your-vertex-endpoint-id';
31+
32+
// Prompt used in the prediction
33+
const prompt = 'Why is the sky blue?';
34+
35+
// Encapsulate the prompt in a correct format for TPUs
36+
// Example format: [{prompt: 'Why is the sky blue?', temperature: 0.9}]
37+
const input = {
38+
prompt,
39+
// Parameters for default configuration
40+
maxOutputTokens: 1024,
41+
temperature: 0.9,
42+
topP: 1.0,
43+
topK: 1,
44+
};
45+
46+
// Convert input message to a list of GAPIC instances for model input
47+
const instances = [helpers.toValue(input)];
48+
49+
// TODO(developer): Uncomment apiEndpoint and predictionServiceClient before running the sample.
50+
// const apiEndpoint = `${endpointRegion}-aiplatform.googleapis.com`;
51+
52+
// Create a client
53+
// predictionServiceClient = new PredictionServiceClient({apiEndpoint});
54+
55+
// Call the Gemma2 endpoint
56+
const gemma2Endpoint = `projects/${projectId}/locations/${endpointRegion}/endpoints/${endpointId}`;
57+
58+
const [response] = await predictionServiceClient.predict({
59+
endpoint: gemma2Endpoint,
60+
instances,
61+
});
62+
63+
const predictions = response.predictions;
64+
const text = predictions[0].stringValue;
65+
66+
console.log('Predictions:', text);
67+
// [END generativeaionvertexai_gemma2_predict_tpu]
68+
return text;
69+
}
70+
71+
module.exports = gemma2PredictTpu;
72+
73+
// TODO(developer): Uncomment below lines before running the sample.
74+
// gemma2PredictTpu(...process.argv.slice(2)).catch(err => {
75+
// console.error(err.message);
76+
// process.exitCode = 1;
77+
// });

‎ai-platform/snippets/test/gemma2Prediction.test.js

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
'use strict';
18+
19+
const {expect} = require('chai');
20+
const {afterEach, describe, it} = require('mocha');
21+
const sinon = require('sinon');
22+
const gemma2PredictGpu = require('../gemma2PredictGpu.js');
23+
const gemma2PredictTpu = require('../gemma2PredictTpu.js');
24+
25+
const gpuResponse = `The sky appears blue due to a phenomenon called **Rayleigh scattering**.
26+
**Here's how it works:**
27+
1. **Sunlight:** Sunlight is composed of all the colors of the rainbow.
28+
2. **Earth's Atmosphere:** When sunlight enters the Earth's atmosphere, it collides with tiny particles like nitrogen and oxygen molecules.
29+
3. **Scattering:** These particles scatter the sunlight in all directions. However, blue light (which has a shorter wavelength) is scattered more effectively than other colors.
30+
4. **Our Perception:** As a result, we see a blue sky because the scattered blue light reaches our eyes from all directions.
31+
**Why not other colors?**
32+
* **Violet light** has an even shorter wavelength than blue and is scattered even more. However, our eyes are less sensitive to violet light, so we perceive the sky as blue.
33+
* **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.
34+
`;
35+
36+
const tpuResponse =
37+
'The sky appears blue due to a phenomenon called **Rayleigh scattering**.';
38+
39+
describe('Gemma2 predictions', async () => {
40+
const gemma2Endpoint =
41+
'projects/your-project-id/locations/your-vertex-endpoint-region/endpoints/your-vertex-endpoint-id';
42+
const configValues = {
43+
maxOutputTokens: {kind: 'numberValue', numberValue: 1024},
44+
temperature: {kind: 'numberValue', numberValue: 0.9},
45+
topP: {kind: 'numberValue', numberValue: 1},
46+
topK: {kind: 'numberValue', numberValue: 1},
47+
};
48+
const prompt = 'Why is the sky blue?';
49+
const predictionServiceClientMock = {
50+
predict: sinon.stub().resolves([]),
51+
};
52+
53+
afterEach(() => {
54+
sinon.reset();
55+
});
56+
57+
it('should run interference with GPU', async () => {
58+
const expectedGpuRequest = {
59+
endpoint: gemma2Endpoint,
60+
instances: [
61+
{
62+
kind: 'structValue',
63+
structValue: {
64+
fields: {
65+
inputs: {
66+
kind: 'stringValue',
67+
stringValue: prompt,
68+
},
69+
parameters: {
70+
kind: 'structValue',
71+
structValue: {
72+
fields: configValues,
73+
},
74+
},
75+
},
76+
},
77+
},
78+
],
79+
};
80+
81+
predictionServiceClientMock.predict.resolves([
82+
{
83+
predictions: [
84+
{
85+
stringValue: gpuResponse,
86+
},
87+
],
88+
},
89+
]);
90+
91+
const output = await gemma2PredictGpu(predictionServiceClientMock);
92+
93+
expect(output).include('Rayleigh scattering');
94+
expect(predictionServiceClientMock.predict.calledOnce).to.be.true;
95+
expect(predictionServiceClientMock.predict.calledWith(expectedGpuRequest))
96+
.to.be.true;
97+
});
98+
99+
it('should run interference with TPU', async () => {
100+
const expectedTpuRequest = {
101+
endpoint: gemma2Endpoint,
102+
instances: [
103+
{
104+
kind: 'structValue',
105+
structValue: {
106+
fields: {
107+
...configValues,
108+
prompt: {
109+
kind: 'stringValue',
110+
stringValue: prompt,
111+
},
112+
},
113+
},
114+
},
115+
],
116+
};
117+
118+
predictionServiceClientMock.predict.resolves([
119+
{
120+
predictions: [
121+
{
122+
stringValue: tpuResponse,
123+
},
124+
],
125+
},
126+
]);
127+
128+
const output = await gemma2PredictTpu(predictionServiceClientMock);
129+
130+
expect(output).include('Rayleigh scattering');
131+
expect(predictionServiceClientMock.predict.calledOnce).to.be.true;
132+
expect(predictionServiceClientMock.predict.calledWith(expectedTpuRequest))
133+
.to.be.true;
134+
});
135+
});

0 commit comments

Comments
 (0)