Fix decoding an empty sequence

elixir-nx · Aug 3, 2023 · 0154888 · 0154888
1 parent 888f53b
commit 0154888
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/lib/bumblebee/utils/tokenizers.ex b/lib/bumblebee/utils/tokenizers.ex
@@ -151,15 +151,15 @@ defmodule Bumblebee.Utils.Tokenizers do
     |> Nx.reshape({length(list), :auto})
   end
 
-  def decode(tokenizer, [id | _] = ids) when is_number(id) do
-    case Tokenizer.decode(tokenizer, ids) do
+  def decode(tokenizer, [ids | _] = batch_ids) when is_list(ids) do
+    case Tokenizer.decode_batch(tokenizer, batch_ids) do
       {:ok, decoded} -> decoded
       {:error, term} -> raise "decoding failed with error: #{inspect(term)}"
     end
   end
 
-  def decode(tokenizer, batch_ids) do
-    case Tokenizer.decode_batch(tokenizer, batch_ids) do
+  def decode(tokenizer, ids) do
+    case Tokenizer.decode(tokenizer, ids) do
       {:ok, decoded} -> decoded
       {:error, term} -> raise "decoding failed with error: #{inspect(term)}"
     end