Skip to content

Commit 64b3da6

Browse files
NikhilVermaxenova
andauthored
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]>
1 parent 250fcd7 commit 64b3da6

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

src/pipelines.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1255,7 +1255,7 @@ export class FeatureExtractionPipeline extends (/** @type {new (options: TextPip
12551255
// console.log(outputs)
12561256

12571257
/** @type {Tensor} */
1258-
let result = outputs.last_hidden_state ?? outputs.logits;
1258+
let result = outputs.last_hidden_state ?? outputs.logits ?? outputs.token_embeddings;
12591259
if (pooling === 'none') {
12601260
// Skip pooling
12611261
} else if (pooling === 'mean') {

0 commit comments

Comments
 (0)