Use lru_cache for tokenization tests (#36818)
* fix * fix * fix * fix --------- Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
@@ -14,12 +14,13 @@
|
||||
# limitations under the License.
|
||||
|
||||
import unittest
|
||||
from functools import lru_cache
|
||||
from typing import Tuple
|
||||
|
||||
from transformers import AddedToken, LukeTokenizer
|
||||
from transformers.testing_utils import get_tests_dir, require_torch, slow
|
||||
|
||||
from ...test_tokenization_common import TokenizerTesterMixin
|
||||
from ...test_tokenization_common import TokenizerTesterMixin, use_cache_if_possible
|
||||
|
||||
|
||||
SAMPLE_VOCAB = get_tests_dir("fixtures/vocab.json")
|
||||
@@ -33,13 +34,17 @@ class LukeTokenizerTest(TokenizerTesterMixin, unittest.TestCase):
|
||||
test_rust_tokenizer = False
|
||||
from_pretrained_kwargs = {"cls_token": "<s>"}
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
|
||||
self.special_tokens_map = {"entity_token_1": "<ent>", "entity_token_2": "<ent2>"}
|
||||
cls.special_tokens_map = {"entity_token_1": "<ent>", "entity_token_2": "<ent2>"}
|
||||
|
||||
def get_tokenizer(self, task=None, **kwargs):
|
||||
kwargs.update(self.special_tokens_map)
|
||||
@classmethod
|
||||
@use_cache_if_possible
|
||||
@lru_cache(maxsize=64)
|
||||
def get_tokenizer(cls, task=None, **kwargs):
|
||||
kwargs.update(cls.special_tokens_map)
|
||||
tokenizer = LukeTokenizer(
|
||||
vocab_file=SAMPLE_VOCAB,
|
||||
merges_file=SAMPLE_MERGE_FILE,
|
||||
@@ -137,8 +142,8 @@ class LukeTokenizerTest(TokenizerTesterMixin, unittest.TestCase):
|
||||
def test_embeded_special_tokens(self):
|
||||
for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
|
||||
with self.subTest("{} ({})".format(tokenizer.__class__.__name__, pretrained_name)):
|
||||
tokenizer_r = self.rust_tokenizer_class.from_pretrained(pretrained_name, **kwargs)
|
||||
tokenizer_p = self.tokenizer_class.from_pretrained(pretrained_name, **kwargs)
|
||||
tokenizer_r = self.get_rust_tokenizer(pretrained_name, **kwargs)
|
||||
tokenizer_p = self.get_tokenizer(pretrained_name, **kwargs)
|
||||
sentence = "A, <mask> AllenNLP sentence."
|
||||
tokens_r = tokenizer_r.encode_plus(sentence, add_special_tokens=True, return_token_type_ids=True)
|
||||
tokens_p = tokenizer_p.encode_plus(sentence, add_special_tokens=True, return_token_type_ids=True)
|
||||
|
||||
Reference in New Issue
Block a user