Allow add_tokens for ESM

src/transformers/models/esm/tokenization_esm.py

@@ -16,8 +16,9 @@
 import os
 from typing import List, Optional, Union
 
+from transformers import AddedToken
+
 from ...tokenization_utils import PreTrainedTokenizer
-from ...tokenization_utils_base import AddedToken
 from ...utils import logging
 
 
@@ -65,8 +66,6 @@ def __init__(
         **kwargs,
     ):
         self.all_tokens = load_vocab_file(vocab_file)
-        self._id_to_token = dict(enumerate(self.all_tokens))
-        self._token_to_id = {tok: ind for ind, tok in enumerate(self.all_tokens)}
         super().__init__(
             unk_token=unk_token,
             cls_token=cls_token,
@@ -82,6 +81,14 @@ def __init__(
         self.unique_no_split_tokens = self.all_tokens
         self._update_trie(self.unique_no_split_tokens)
 
+    @property
+    def _id_to_token(self):
+        return dict(enumerate(self.all_tokens))
+
+    @property
+    def _token_to_id(self):
+        return {token: i for i, token in enumerate(self.all_tokens)}
+
     def _convert_id_to_token(self, index: int) -> str:
         return self._id_to_token.get(index, self.unk_token)
 
@@ -159,4 +166,8 @@ def vocab_size(self) -> int:
         return self.get_vocab_size(with_added_tokens=False)
 
     def _add_tokens(self, new_tokens: Union[List[str], List[AddedToken]], special_tokens: bool = False) -> int:
-        return super()._add_tokens(new_tokens, special_tokens=True)
+        result = super()._add_tokens(new_tokens, special_tokens=special_tokens)
+        for token in new_tokens:
+            if token and token not in self.all_tokens:
+                self.all_tokens.append(token)
+        return result

tests/models/esm/test_tokenization_esm.py

@@ -87,3 +87,25 @@ def test_tokenize_special_tokens(self):
                 self.assertEqual(len(token_2), 1)
                 self.assertEqual(token_1[0], SPECIAL_TOKEN_1)
                 self.assertEqual(token_2[0], SPECIAL_TOKEN_2)
+
+    def test_add_tokens(self):
+        tokenizer = self.tokenizer_class(self.vocab_file)
+
+        vocab_size = len(tokenizer)
+        self.assertEqual(tokenizer.add_tokens(""), 0)
+        self.assertEqual(tokenizer.add_tokens("testoken"), 1)
+        self.assertEqual(tokenizer.add_tokens(["testoken1", "testtoken2"]), 2)
+        self.assertEqual(len(tokenizer), vocab_size + 3)
+
+        self.assertEqual(tokenizer.add_special_tokens({}), 0)
+        self.assertEqual(tokenizer.add_special_tokens({"bos_token": "[BOS]", "eos_token": "[EOS]"}), 2)
+        self.assertRaises(AssertionError, tokenizer.add_special_tokens, {"additional_special_tokens": "<testtoken1>"})
+        self.assertEqual(tokenizer.add_special_tokens({"additional_special_tokens": ["<testtoken2>"]}), 1)
+        self.assertEqual(
+            tokenizer.add_special_tokens({"additional_special_tokens": ["<testtoken3>", "<testtoken4>"]}), 2
+        )
+        self.assertIn("<testtoken3>", tokenizer.special_tokens_map["additional_special_tokens"])
+        self.assertIsInstance(tokenizer.special_tokens_map["additional_special_tokens"], list)
+        self.assertGreaterEqual(len(tokenizer.special_tokens_map["additional_special_tokens"]), 2)
+
+        self.assertEqual(len(tokenizer), vocab_size + 8)