Fix batched text generation finishing too early

elixir-nx · Jul 4, 2024 · 3b0eb08 · 3b0eb08
1 parent 6cf6839
commit 3b0eb08
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/lib/bumblebee/text/generation.ex b/lib/bumblebee/text/generation.ex
@@ -609,8 +609,10 @@ defmodule Bumblebee.Text.Generation do
       eos_token_ids = List.wrap(eos_token_id)
 
       token_id
+      |> Nx.vectorize(:batch)
       |> Nx.equal(Nx.tensor(eos_token_ids))
       |> Nx.any()
+      |> Nx.devectorize()
     else
       Nx.tensor(false)
     end

diff --git a/test/bumblebee/audio/speech_to_text_whisper_test.exs b/test/bumblebee/audio/speech_to_text_whisper_test.exs
@@ -109,7 +109,7 @@ defmodule Bumblebee.Audio.SpeechToTextWhisperTest do
         generation_config,
         chunk_num_seconds: 30,
         defn_options: [compiler: EXLA],
-        compile: [batch_size: 1]
+        compile: [batch_size: 4]
       )
 
     audio =