Add CLS pooling option to feature-extraction pipeline (#450)

huggingface · Dec 11, 2023 · 2cd2997 · 2cd2997
1 parent 8e49e5e
commit 2cd2997
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/pipelines.js b/src/pipelines.js
@@ -933,6 +933,8 @@ export class FeatureExtractionPipeline extends Pipeline {
             // Skip pooling
         } else if (pooling === 'mean') {
             result = mean_pooling(result, inputs.attention_mask);
+        } else if (pooling === 'cls') {
+            result = result.slice(null, 0);
         } else {
             throw Error(`Pooling method '${pooling}' not supported.`);
         }