Update public-apps/whisper-chat-app.livemd

public-apps/whisper-chat-app.livemd

@@ -22,16 +22,21 @@ This chat is open to anyone, be polite and act responsibly. :) Note chat history
 
 ```elixir
 {:ok, model_info} =
-  Bumblebee.load_model({:hf, "openai/whisper-tiny"}, log_params_diff: false)
+  Bumblebee.load_model({:hf, "openai/whisper-small"}, log_params_diff: false)
 
-{:ok, featurizer} = Bumblebee.load_featurizer({:hf, "openai/whisper-tiny"})
-{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/whisper-tiny"})
+{:ok, featurizer} = Bumblebee.load_featurizer({:hf, "openai/whisper-small"})
+{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/whisper-small"})
 
 serving =
   Bumblebee.Audio.speech_to_text(model_info, featurizer, tokenizer,
     max_new_tokens: 100,
     compile: [batch_size: 8],
-    defn_options: [compiler: EXLA]
+    defn_options: [compiler: EXLA],
+    forced_token_ids: [
+      {1, Bumblebee.Tokenizer.token_to_id(tokenizer, "<|pt|>")},
+      {2, Bumblebee.Tokenizer.token_to_id(tokenizer, "<|transcribe|>")},
+      {3, Bumblebee.Tokenizer.token_to_id(tokenizer, "<|notimestamps|>")}
+    ]
   )
 
 Kino.start_child({Nx.Serving, serving: serving, name: WhisperChat})
@@ -51,8 +56,13 @@ Kino.async_listen(form, fn %{data: %{audio: audio, name: name}, origin: origin}
       |> Nx.reshape({:auto, audio.num_channels})
       |> Nx.mean(axes: [1])
 
-    %{results: [%{text: generated_text}]} = Nx.Serving.batched_run(WhisperChat, audio)
-    content = Kino.Markdown.new("**#{name}**: #{generated_text}")
+    {time, your_func_result} = :timer.tc(&Nx.Serving.batched_run/2, [WhisperChat, audio])
+
+    %{results: [%{text: generated_text}]} = your_func_result
+
+    content =
+      Kino.Markdown.new("**#{name}** (in #{time / (1000 * 1000)} seconds) : #{generated_text}")
+
     Kino.Frame.append(frame, content, temporary: true)
   else
     content = Kino.Markdown.new("*Error! Name and Audio are required*")