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.
xenova · May 22, 2024 · c87a0b9 · c87a0b9
1 parent 0af1e2f
commit c87a0b9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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;
-}
+}