Fix loading Gemma tokenizer

elixir-nx · Mar 14, 2024 · 21a2533 · 21a2533
1 parent 93b5580
commit 21a2533
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 4 deletions.
diff --git a/lib/bumblebee/text/pre_trained_tokenizer.ex b/lib/bumblebee/text/pre_trained_tokenizer.ex
@@ -141,10 +141,12 @@ defmodule Bumblebee.Text.PreTrainedTokenizer do
       special_tokens: %{unk: "[UNK]", sep: "[SEP]", pad: "[PAD]", cls: "[CLS]", mask: "[MASK]"}
     },
     gemma: %{
-      unk: "<unk>",
-      bos: "<bos>",
-      eos: "<eos>",
-      pad: "<pad>"
+      special_tokens: %{
+        unk: "<unk>",
+        bos: "<bos>",
+        eos: "<eos>",
+        pad: "<pad>"
+      }
     },
     gpt_neo_x: %{
       special_tokens: %{

diff --git a/test/bumblebee/text/pre_trained_tokenizer_test.exs b/test/bumblebee/text/pre_trained_tokenizer_test.exs
@@ -211,6 +211,42 @@ defmodule Bumblebee.Text.PreTrainedTokenizerTest do
     )
   end
 
+  test ":gemma" do
+    assert {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "unsloth/gemma-7b-it"})
+
+    assert %Bumblebee.Text.PreTrainedTokenizer{type: :gemma} = tokenizer
+
+    inputs = Bumblebee.apply_tokenizer(tokenizer, ["Hello World"])
+
+    assert_equal(
+      inputs["input_ids"],
+      Nx.tensor([[2, 4521, 3855]])
+    )
+
+    assert_equal(
+      inputs["attention_mask"],
+      Nx.tensor([[1, 1, 1]])
+    )
+  end
+
+  test ":gpt_neo_x" do
+    assert {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "EleutherAI/gpt-neox-20b"})
+
+    assert %Bumblebee.Text.PreTrainedTokenizer{type: :gpt_neo_x} = tokenizer
+
+    inputs = Bumblebee.apply_tokenizer(tokenizer, ["Hello World"])
+
+    assert_equal(
+      inputs["input_ids"],
+      Nx.tensor([[12092, 3645]])
+    )
+
+    assert_equal(
+      inputs["attention_mask"],
+      Nx.tensor([[1, 1]])
+    )
+  end
+
   test ":gpt2" do
     assert {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai-community/gpt2"})