feat: Elastic Text-Embedding Model demo. (#3650)

Author: lee1premium-2024

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

@@ -0,0 +1,62 @@
+/*
+ * 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';
+
+// [START generativeaionvertexai_sdk_embedding]
+async function main(
+  project,
+  model = 'text-embedding-preview-0409',
+  texts = 'banana bread?;banana muffins?',
+  task = 'QUESTION_ANSWERING',
+  outputDimensionality = 256,
+  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 match = apiEndpoint.match(/(?<Location>\w+-\w+)/);
+  const location = match ? match.groups.Location : 'us-centra11';
+  const endpoint = `projects/${project}/locations/${location}/publishers/google/models/${model}`;
+  const parameters = helpers.toValue(outputDimensionality);
+
+  async function callPredict() {
+    const instances = texts
+      .split(';')
+      .map(e => helpers.toValue({content: e, taskType: task}));
+    const request = {endpoint, instances, parameters};
+    const client = new PredictionServiceClient(clientOptions);
+    const [response] = await client.predict(request);
+    console.log('Got predict response');
+    const predictions = response.predictions;
+    for (const prediction of predictions) {
+      const embeddings = prediction.structValue.fields.embeddings;
+      const values = embeddings.structValue.fields.values.listValue.values;
+      console.log('Got prediction: ' + JSON.stringify(values));
+    }
+  }
+
+  callPredict();
+}
+// [END generativeaionvertexai_sdk_embedding]
+
+process.on('unhandledRejection', err => {
+  console.error(err.message);
+  process.exitCode = 1;
+});
+
+main(...process.argv.slice(2));

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

@@ -16,70 +16,41 @@
 
 'use strict';
 
-async function main(project, location = 'us-central1') {
-  // [START aiplatform_sdk_embedding]
-  /**
-   * 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';
+// [START aiplatform_sdk_embedding]
+async function main(
+  project,
+  model = 'textembedding-gecko@003',
+  texts = 'banana bread?;banana muffins?',
+  task = 'RETRIEVAL_DOCUMENT',
+  apiEndpoint = 'us-central1-aiplatform.googleapis.com'
+) {
   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 = 'textembedding-gecko@001';
-
-  // Instantiates a client
-  const predictionServiceClient = new PredictionServiceClient(clientOptions);
+  const {helpers} = aiplatform; // helps construct protobuf.Value objects.
+  const clientOptions = {apiEndpoint: apiEndpoint};
+  const match = apiEndpoint.match(/(?<Location>\w+-\w+)/);
+  const location = match ? match.groups.Location : 'us-centra11';
+  const endpoint = `projects/${project}/locations/${location}/publishers/google/models/${model}`;
 
   async function callPredict() {
-    // Configure the parent resource
-    const endpoint = `projects/${project}/locations/${location}/publishers/${publisher}/models/${model}`;
-
-    const instance = {
-      content: 'What is life?',
-    };
-    const instanceValue = helpers.toValue(instance);
-    const instances = [instanceValue];
-
-    const parameter = {
-      temperature: 0,
-      maxOutputTokens: 256,
-      topP: 0,
-      topK: 1,
-    };
-    const parameters = helpers.toValue(parameter);
-
-    const request = {
-      endpoint,
-      instances,
-      parameters,
-    };
-
-    // Predict request
-    const [response] = await predictionServiceClient.predict(request);
-    console.log('Get text embeddings response');
+    const instances = texts
+      .split(';')
+      .map(e => helpers.toValue({content: e, taskType: task}));
+    const request = {endpoint, instances};
+    const client = new PredictionServiceClient(clientOptions);
+    const [response] = await client.predict(request);
+    console.log('Got predict response');
     const predictions = response.predictions;
-    console.log('\tPredictions :');
     for (const prediction of predictions) {
-      console.log(`\t\tPrediction : ${JSON.stringify(prediction)}`);
+      const embeddings = prediction.structValue.fields.embeddings;
+      const values = embeddings.structValue.fields.values.listValue.values;
+      console.log('Got prediction: ' + JSON.stringify(values));
     }
   }
 
   callPredict();
-  // [END aiplatform_sdk_embedding]
 }
+// [END aiplatform_sdk_embedding]
 
 process.on('unhandledRejection', err => {
   console.error(err.message);

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

@@ -1,14 +1,14 @@
 /*
  * Copyright 2023 Google LLC
  *
- * Licensed under the Apache License, Version 2.0 (the "License");
+ * 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,
+ * 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.
@@ -25,16 +25,27 @@ 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';
+const texts = [
+  'banana bread?',
+  'banana muffin?',
+  'banana?',
+  'recipe?',
+  'muffin recipe?',
+].join(';');
 
-describe('AI platform predict text embeddings', () => {
-  it('should make predictions using a large language model', async () => {
+describe('predict text embeddings', () => {
+  it('should get text embeddings using the latest model', async () => {
     const stdout = execSync(
-      `node ./predict-text-embeddings.js ${project} ${location}`,
-      {
-        cwd,
-      }
+      `node ./predict-text-embeddings.js ${project} textembedding-gecko@003 '${texts}' RETRIEVAL_DOCUMENT`,
+      {cwd}
     );
-    assert.match(stdout, /Get text embeddings response/);
+    assert.match(stdout, /Got predict response/);
+  });
+  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}' QUESTION_ANSWERING 256`,
+      {cwd}
+    );
+    assert.match(stdout, /Got predict response/);
   });
 });