From 2cd2997d5746f89b1a1dadd80a6fb85827236d42 Mon Sep 17 00:00:00 2001
From: Joshua Lochner <admin@xenova.com>
Date: Tue, 12 Dec 2023 00:19:59 +0200
Subject: [PATCH] Add CLS pooling option to `feature-extraction` pipeline
 (#450)

---
 src/pipelines.js | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/pipelines.js b/src/pipelines.js
index 231ea704e..af4594b53 100644
--- 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.`);
         }