Skip to content

Commit

Permalink
Merge pull request #3262 from GoogleCloudPlatform/b279595025-unittest
Browse files Browse the repository at this point in the history
feat(aiplatform): code generation unit test sample for Vertex LLMs
  • Loading branch information
gericdong authored Jun 7, 2023
2 parents 0bd9bc6 + 3ee85a3 commit a9c82de
Show file tree
Hide file tree
Showing 2 changed files with 138 additions and 0 deletions.
98 changes: 98 additions & 0 deletions ai-platform/snippets/predict-code-generation-unittest.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,98 @@
/*
* Copyright 2023 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.
*/

'use strict';

async function main(project, location = 'us-central1') {
// [START aiplatform_sdk_code_generation_unittest]
/**
* TODO(developer): Uncomment these variables before running the sample.\
* (Not necessary if passing values as arguments)
*/
// const project = 'YOUR_PROJECT_ID';
// const location = 'YOUR_PROJECT_LOCATION';
const aiplatform = require('@google-cloud/aiplatform');

// Imports the Google Cloud Prediction service client
const {PredictionServiceClient} = aiplatform.v1;

// Import the helper module for converting arbitrary protobuf.Value objects.
const {helpers} = aiplatform;

// Specifies the location of the api endpoint
const clientOptions = {
apiEndpoint: 'us-central1-aiplatform.googleapis.com',
};
const publisher = 'google';
const model = 'code-bison@001';

// Instantiates a client
const predictionServiceClient = new PredictionServiceClient(clientOptions);

async function callPredict() {
// Configure the parent resource
const endpoint = `projects/${project}/locations/${location}/publishers/${publisher}/models/${model}`;

const prompt = {
prefix:
'Write a unit test for this function: \
def is_leap_year(year): \
if year % 4 == 0: \
if year % 100 == 0: \
if year % 400 == 0: \
return True \
else: \
return False \
else: \
return True \
else: \
return False',
};
const instanceValue = helpers.toValue(prompt);
const instances = [instanceValue];

const parameter = {
temperature: 0.2,
maxOutputTokens: 512,
};
const parameters = helpers.toValue(parameter);

const request = {
endpoint,
instances,
parameters,
};

// Predict request
const [response] = await predictionServiceClient.predict(request);
console.log('Get code generation response');
const predictions = response.predictions;
console.log('\tPredictions :');
for (const prediction of predictions) {
console.log(`\t\tPrediction : ${JSON.stringify(prediction)}`);
}
}

callPredict();
// [END aiplatform_sdk_code_generation_unittest]
}

process.on('unhandledRejection', err => {
console.error(err.message);
process.exitCode = 1;
});

main(...process.argv.slice(2));
40 changes: 40 additions & 0 deletions ai-platform/snippets/test/predict-code-generation-unittest.test.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/*
* Copyright 2023 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.
*/

'use strict';

const path = require('path');
const {assert} = require('chai');
const {describe, it} = require('mocha');

const cp = require('child_process');
const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'});
const cwd = path.join(__dirname, '..');

const project = process.env.CAIP_PROJECT_ID;
const location = 'us-central1';

describe('AI platform predict code generation', () => {
it('should make predictions using a large language model', async () => {
const stdout = execSync(
`node ./predict-code-generation-unittest.js ${project} ${location}`,
{
cwd,
}
);
assert.match(stdout, /Get code generation response/);
});
});

0 comments on commit a9c82de

Please sign in to comment.