add a small test

huggingface · Oct 4, 2024 · ed34ffd · ed34ffd
1 parent b5640a6
commit ed34ffd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/bindings/python/py_src/tokenizers/__init__.pyi b/bindings/python/py_src/tokenizers/__init__.pyi
@@ -725,7 +725,7 @@ class Tokenizer:
         """
         pass
 
-    def assing_tokens(self, old_tokens, new_tokens):
+    def assign_tokens(self, old_tokens, new_tokens):
         """
         Add the given tokens to the vocabulary
 

diff --git a/bindings/python/tests/bindings/test_tokenizer.py b/bindings/python/tests/bindings/test_tokenizer.py
@@ -562,6 +562,13 @@ def test_setting_to_none(self):
         tokenizer.pre_tokenizer = None
         assert tokenizer.pre_tokenizer == None
 
+    def test_re_assign_tokens(self):
+        tokenizer = Tokenizer.from_pretrained("t5-base")
+        tokenizer.assign_tokens({"<extra_id_0>": "my_new_token"})
+        assert tokenizer.decode([32099]) == "my_new_token"
+        assert tokenizer.encode("<extra_id_0>").tokens == ["▁", "<", "extra", "_", "i", "d", "_", "0", ">", "</s>"]
+        assert "my_new_token" in tokenizer.get_vocab(True).keys()
+
 
 class TestTokenizerRepr:
     def test_repr(self):