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 <[email protected]>
huggingface · May 23, 2024 · 64b3da6 · 64b3da6
1 parent 250fcd7
commit 64b3da6
Showing 1 changed file with 1 addition and 1 deletion.
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') {