Merge branch 'main' into batch-create-custom-network

Author: subfuzion

ai-platform/snippets/predict-text-embeddings-preview.js

@@ -16,20 +16,25 @@
 
 'use strict';
 
-// [START generativeaionvertexai_sdk_embedding]
-async function main(
-  project,
-  model = 'text-embedding-preview-0409',
-  texts = 'banana bread?;banana muffins?',
-  task = 'QUESTION_ANSWERING',
-  dimensionality = 256,
-  apiEndpoint = 'us-central1-aiplatform.googleapis.com'
-) {
+async function main() {
+  // [START generativeaionvertexai_sdk_embedding]
+
+  // TODO(developer): Update the following for your own use case.
+  const project = 'long-door-651';
+  const model = 'text-embedding-preview-0815';
+  const location = 'us-central1';
+  // Calculate the embedding for code blocks. Using 'RETRIEVAL_DOCUMENT' for corpus.
+  // Specify the task type as 'CODE_RETRIEVAL_QUERY' for query, e.g. 'Retrieve a function that adds two numbers'.
+  const texts =
+    'def func(a, b): return a + b;def func(a, b): return a - b;def func(a, b): return (a ** 2 + b ** 2) ** 0.5';
+  const task = 'RETRIEVAL_DOCUMENT';
+  const dimensionality = 3;
+  const apiEndpoint = 'us-central1-aiplatform.googleapis.com';
+
   const aiplatform = require('@google-cloud/aiplatform');
   const {PredictionServiceClient} = aiplatform.v1;
   const {helpers} = aiplatform; // helps construct protobuf.Value objects.
   const clientOptions = {apiEndpoint: apiEndpoint};
-  const location = 'us-central1';
   const endpoint = `projects/${project}/locations/${location}/publishers/google/models/${model}`;
   const parameters = helpers.toValue({
     outputDimensionality: parseInt(dimensionality),
@@ -50,14 +55,11 @@ async function main(
     });
     console.log('Got embeddings: \n' + JSON.stringify(embeddings));
   }
-
-  callPredict();
+  await callPredict();
+  // [END generativeaionvertexai_sdk_embedding]
 }
-// [END generativeaionvertexai_sdk_embedding]
 
-process.on('unhandledRejection', err => {
-  console.error(err.message);
+main().catch(err => {
+  console.error(err);
   process.exitCode = 1;
 });
-
-main(...process.argv.slice(2));

ai-platform/snippets/test/predict-text-embeddings.test.js

@@ -47,12 +47,9 @@ describe('predict text embeddings', () => {
     }
   });
   it('should get text embeddings using the preview model', async () => {
-    const stdout = execSync(
-      `node ./predict-text-embeddings-preview.js ${project} text-embedding-preview-0409 '${texts.join(';')}' QUESTION_ANSWERING ${dimensionality}`,
-      {cwd}
-    );
+    const stdout = execSync('node ./predict-text-embeddings-preview.js', {cwd});
     const embeddings = JSON.parse(stdout.trimEnd().split('\n').at(-1));
-    assert.equal(texts.length, embeddings.length);
+    assert.equal(3, embeddings.length);
     for (const embedding of embeddings) {
       assert.equal(dimensionality, embedding.length);
     }