Check text embedding output shape before applying pooling

elixir-nx · Dec 18, 2023 · b374fa6 · b374fa6
1 parent 6a139fa
commit b374fa6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/lib/bumblebee/text.ex b/lib/bumblebee/text.ex
@@ -302,8 +302,8 @@ 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`. By
-      default no pooling is applied
+      it is not already a pooled embedding. Supported values: `:mean_pooling`.
+      By default no pooling is applied
 
     * `:embedding_processor` - a post-processing step to apply to the
       embedding. Supported values: `:l2_norm`. By default the output is

diff --git a/lib/bumblebee/text/text_embedding.ex b/lib/bumblebee/text/text_embedding.ex
@@ -62,6 +62,16 @@ defmodule Bumblebee.Text.TextEmbedding do
             output
 
           :mean_pooling ->
+            case Nx.rank(output) do
+              3 ->
+                :ok
+
+              rank ->
+                raise ArgumentError,
+                      "expected the output tensor to have rank 3 with :output_pool is enabled, got: #{rank}." <>
+                        " You should either disable pooling or pick a different output using :output_attribute"
+            end
+
             input_mask_expanded = Nx.new_axis(inputs["attention_mask"], -1)
 
             output