Improve error message in text embedding when configured output is Axon.None

Author: jonatanklosko

lib/bumblebee/text.ex

@@ -381,7 +381,7 @@ defmodule Bumblebee.Text do
 
         * `:mean_pooling` - performs a mean across all tokens
 
-        * `cls_token_pooling` - takes the embedding for the special CLS token.
+        * `: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
 

lib/bumblebee/text/text_embedding.ex

@@ -42,6 +42,14 @@ defmodule Bumblebee.Text.TextEmbedding do
 
       output =
         case output do
+          %{^output_attribute => %Axon.None{}} ->
+            keys = output |> Map.keys() |> Enum.sort()
+
+            raise ArgumentError,
+                  "key #{inspect(output_attribute)} in the output map has value %Axon.None{}," <>
+                    " make sure this is the correct key and check module documentation incase it is opt-in." <>
+                    " The output map keys are: #{inspect(keys)}"
+
           %{^output_attribute => output} ->
             output