From 64b3da608bf35247d19082824e5a1119063bbed9 Mon Sep 17 00:00:00 2001 From: Nikhil Verma Date: Fri, 24 May 2024 02:15:32 +0530 Subject: [PATCH] Update pipelines.js to allow for `token_embeddings` as well (#770) * 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. * Keep diff small * Keep diff small --------- Co-authored-by: Joshua Lochner --- src/pipelines.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pipelines.js b/src/pipelines.js index 805f64736..c7772aa55 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') {