Skip to content

Commit

Permalink
Update pipelines.js to allow for token_embeddings as well (#770)
Browse files Browse the repository at this point in the history
* 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 <admin@xenova.com>
  • Loading branch information
NikhilVerma and xenova committed May 23, 2024
1 parent 250fcd7 commit 64b3da6
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion src/pipelines.js
Original file line number Diff line number Diff line change
Expand Up @@ -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') {
Expand Down

0 comments on commit 64b3da6

Please sign in to comment.