From 7db36b821f0a2a60c1d1ee33ae498ad51f38bc6c Mon Sep 17 00:00:00 2001
From: Niko Maroulis <nik.maroulis@gmail.com>
Date: Tue, 6 Aug 2024 09:12:11 -0400
Subject: [PATCH] Add support for CLS token pooling in text embedding (#385)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>
---
 lib/bumblebee/text.ex                |  9 ++++++++-
 lib/bumblebee/text/text_embedding.ex | 21 ++++++++++-----------
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/lib/bumblebee/text.ex b/lib/bumblebee/text.ex
index 2e9233b0..aa4ca19b 100644
--- a/lib/bumblebee/text.ex
+++ b/lib/bumblebee/text.ex
@@ -313,7 +313,14 @@ defmodule Bumblebee.Text do
       this option is ignored. Defaults to `:pooled_state`
 
     * `:output_pool` - pooling to apply on top of the model output, in case
-      it is not already a pooled embedding. Supported values: `:mean_pooling`.
+      it is not already a pooled embedding. Supported values:
+
+        * `:mean_pooling` - performs a mean across all tokens
+        
+        * `cls_token_pooling` - takes the embedding for the special CLS token.
+          Note that we currently assume that the CLS token is the first token
+          in the sequence
+
       By default no pooling is applied
 
     * `:embedding_processor` - a post-processing step to apply to the
diff --git a/lib/bumblebee/text/text_embedding.ex b/lib/bumblebee/text/text_embedding.ex
index bf6a7a97..41e284e8 100644
--- a/lib/bumblebee/text/text_embedding.ex
+++ b/lib/bumblebee/text/text_embedding.ex
@@ -56,22 +56,21 @@ defmodule Bumblebee.Text.TextEmbedding do
             output
         end
 
+      if output_pool != nil and Nx.rank(output) != 3 do
+        raise ArgumentError,
+              "expected the output tensor to have rank 3 to apply :output_pool, got: #{Nx.rank(output)}." <>
+                " You should either disable pooling or pick a different output using :output_attribute"
+      end
+
       output =
         case output_pool do
           nil ->
             output
 
-          :mean_pooling ->
-            case Nx.rank(output) do
-              3 ->
-                :ok
-
-              rank ->
-                raise ArgumentError,
-                      "expected the output tensor to have rank 3 to apply :output_pool, got: #{rank}." <>
-                        " You should either disable pooling or pick a different output using :output_attribute"
-            end
+          :cls_token_pooling ->
+            Nx.take(output, 0, axis: 1)
 
+          :mean_pooling ->
             input_mask_expanded = Nx.new_axis(inputs["attention_mask"], -1)
 
             output
@@ -81,7 +80,7 @@ defmodule Bumblebee.Text.TextEmbedding do
 
           other ->
             raise ArgumentError,
-                  "expected :output_pool to be one of nil or :mean_pooling, got: #{inspect(other)}"
+                  "expected :output_pool to be one of :cls_token_pooling, :mean_pooling or nil, got: #{inspect(other)}"
         end
 
       output =