more tests

mosaicml · Apr 2, 2024 · cc60a21 · cc60a21
1 parent c6e9f05
commit cc60a21
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/tests/tokenizers/test_tiktoken.py b/tests/tokenizers/test_tiktoken.py
@@ -372,6 +372,15 @@ def test_additional_special_tokens_len():
     no_special = TiktokenTokenizerWrapper(model_name='gpt-4',)
     assert len(with_special.get_vocab()) == len(no_special.get_vocab()) + 1
 
+    ret = with_special.add_special_tokens(
+        {'additional_special_tokens': ['<|im_start|>']})
+    assert ret == 0
+
+    ret = with_special.add_special_tokens(
+        {'additional_special_tokens': ['<|im_end|>']})
+    assert ret == 1
+    assert len(with_special.get_vocab()) == len(no_special.get_vocab()) + 2
+
 
 @pytest.mark.parametrize('model_name,encoding_name',
                          MODEL_ENCODING_NAME_PARAMETRIZATION)