From c87a0b95df8bc8ba326ca7cf4289d0b1077bfa6f Mon Sep 17 00:00:00 2001 From: Nikhil Verma Date: Wed, 22 May 2024 15:49:48 +0530 Subject: [PATCH 1/3] Update pipelines.js to allow for `token_embeddings` as well In recent examples of optimum pipeline export the feature extraction pipelines have their output state as `token_embeddings` instead of `last_hidden_state`. So we should support this as well. --- src/pipelines.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pipelines.js b/src/pipelines.js index 421053b9c..9761359b0 100644 --- a/src/pipelines.js +++ b/src/pipelines.js @@ -1255,7 +1255,7 @@ export class FeatureExtractionPipeline extends (/** @type {new (options: TextPip // console.log(outputs) /** @type {Tensor} */ - let result = outputs.last_hidden_state ?? outputs.logits; + let result = outputs.last_hidden_state ?? outputs.logits ?? outputs.token_embeddings; if (pooling === 'none') { // Skip pooling } else if (pooling === 'mean') { @@ -3282,4 +3282,4 @@ async function loadItems(mapping, model, pretrainedOptions) { } return result; -} \ No newline at end of file +} From 26f1fe0972fb59a5867fddf52761fb1248f80990 Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Thu, 23 May 2024 12:23:32 +0200 Subject: [PATCH 2/3] Keep diff small --- src/pipelines.js | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pipelines.js b/src/pipelines.js index 9761359b0..71defdfb7 100644 --- a/src/pipelines.js +++ b/src/pipelines.js @@ -3283,3 +3283,4 @@ async function loadItems(mapping, model, pretrainedOptions) { return result; } + From 643338be68459e0cd7559029d2c7736b496d9e4e Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Thu, 23 May 2024 10:28:30 +0000 Subject: [PATCH 3/3] Keep diff small --- src/pipelines.js | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/pipelines.js b/src/pipelines.js index 71defdfb7..11a52efd3 100644 --- a/src/pipelines.js +++ b/src/pipelines.js @@ -3282,5 +3282,4 @@ async function loadItems(mapping, model, pretrainedOptions) { } return result; -} - +} \ No newline at end of file