Apply suggestions from code review

elixir-nx · Aug 19, 2024 · 1229a9d · 1229a9d
1 parent 86c2436
commit 1229a9d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex
@@ -255,7 +255,12 @@ defmodule Bumblebee do
     "t5" => :t5,
     "whisper" => :whisper,
     "xlm-roberta" => :xlm_roberta,
-    "m2m_100" => :m2m_100
+    # Both M2M100 and NLLB model checkpoints use the M2M100 model,
+    # but have distinct tokenizers. Consequently, model type is
+    # "m2m_100" in both cases. Currently only NLLB has fast tokenizer
+    # implementation, so if we load the tokenizer correctly, it is
+    # safe to assume it's NLLB
+    "m2m_100" => :nllb
   }
 
   @diffusers_class_to_scheduler %{

diff --git a/lib/bumblebee/text/pre_trained_tokenizer.ex b/lib/bumblebee/text/pre_trained_tokenizer.ex
@@ -179,7 +179,7 @@ defmodule Bumblebee.Text.PreTrainedTokenizer do
         mask: "<mask>"
       }
     },
-    m2m_100: %{
+    nllb: %{
       special_tokens: %{
         eos: "</s>",
         unk: "<unk>",