Use lru_cache for tokenization tests (#36818)

* fix

* fix

* fix

* fix

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

tests/models/roberta/test_tokenization_roberta.py

@@ -18,12 +18,13 @@ import itertools
 import json
 import os
 import unittest
+from functools import lru_cache
 
 from transformers import AddedToken, RobertaTokenizer, RobertaTokenizerFast
 from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES
 from transformers.testing_utils import require_tokenizers, slow
 
-from ...test_tokenization_common import TokenizerTesterMixin
+from ...test_tokenization_common import TokenizerTesterMixin, use_cache_if_possible
 
 
 @require_tokenizers
@@ -34,8 +35,9 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
     test_rust_tokenizer = True
     from_pretrained_kwargs = {"cls_token": "<s>"}
 
-    def setUp(self):
-        super().setUp()
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
 
         # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
         vocab = [
@@ -62,22 +64,30 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
         ]
         vocab_tokens = dict(zip(vocab, range(len(vocab))))
         merges = ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
-        self.special_tokens_map = {"unk_token": "<unk>"}
+        cls.special_tokens_map = {"unk_token": "<unk>"}
 
-        self.vocab_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
-        self.merges_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["merges_file"])
-        with open(self.vocab_file, "w", encoding="utf-8") as fp:
+        cls.vocab_file = os.path.join(cls.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
+        cls.merges_file = os.path.join(cls.tmpdirname, VOCAB_FILES_NAMES["merges_file"])
+        with open(cls.vocab_file, "w", encoding="utf-8") as fp:
             fp.write(json.dumps(vocab_tokens) + "\n")
-        with open(self.merges_file, "w", encoding="utf-8") as fp:
+        with open(cls.merges_file, "w", encoding="utf-8") as fp:
             fp.write("\n".join(merges))
 
-    def get_tokenizer(self, **kwargs):
-        kwargs.update(self.special_tokens_map)
-        return self.tokenizer_class.from_pretrained(self.tmpdirname, **kwargs)
+    @classmethod
+    @use_cache_if_possible
+    @lru_cache(maxsize=64)
+    def get_tokenizer(cls, pretrained_name=None, **kwargs):
+        kwargs.update(cls.special_tokens_map)
+        pretrained_name = pretrained_name or cls.tmpdirname
+        return cls.tokenizer_class.from_pretrained(pretrained_name, **kwargs)
 
-    def get_rust_tokenizer(self, **kwargs):
-        kwargs.update(self.special_tokens_map)
-        return self.rust_tokenizer_class.from_pretrained(self.tmpdirname, **kwargs)
+    @classmethod
+    @use_cache_if_possible
+    @lru_cache(maxsize=64)
+    def get_rust_tokenizer(cls, pretrained_name=None, **kwargs):
+        kwargs.update(cls.special_tokens_map)
+        pretrained_name = pretrained_name or cls.tmpdirname
+        return cls.rust_tokenizer_class.from_pretrained(pretrained_name, **kwargs)
 
     def get_input_output_texts(self, tokenizer):
         input_text = "lower newer"
@@ -171,8 +181,8 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
     def test_embeded_special_tokens(self):
         for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
             with self.subTest(f"{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)
@@ -202,7 +212,7 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
 
     def test_change_add_prefix_space_and_trim_offsets_args(self):
         for trim_offsets, add_prefix_space in itertools.product([True, False], repeat=2):
-            tokenizer_r = self.rust_tokenizer_class.from_pretrained(
+            tokenizer_r = self.get_rust_tokenizer(
                 self.tmpdirname, use_fast=True, add_prefix_space=add_prefix_space, trim_offsets=trim_offsets
             )
 
@@ -222,7 +232,7 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
                 text_of_1_token = "hello"  # `hello` is a token in the vocabulary of `pretrained_name`
                 text = f"{text_of_1_token} {text_of_1_token}"
 
-                tokenizer_r = self.rust_tokenizer_class.from_pretrained(
+                tokenizer_r = self.get_rust_tokenizer(
                     pretrained_name, use_fast=True, add_prefix_space=True, trim_offsets=True
                 )
                 encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
@@ -232,7 +242,7 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
                     (len(text_of_1_token) + 1, len(text_of_1_token) + 1 + len(text_of_1_token)),
                 )
 
-                tokenizer_r = self.rust_tokenizer_class.from_pretrained(
+                tokenizer_r = self.get_rust_tokenizer(
                     pretrained_name, use_fast=True, add_prefix_space=False, trim_offsets=True
                 )
                 encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
@@ -242,7 +252,7 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
                     (len(text_of_1_token) + 1, len(text_of_1_token) + 1 + len(text_of_1_token)),
                 )
 
-                tokenizer_r = self.rust_tokenizer_class.from_pretrained(
+                tokenizer_r = self.get_rust_tokenizer(
                     pretrained_name, use_fast=True, add_prefix_space=True, trim_offsets=False
                 )
                 encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
@@ -252,7 +262,7 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
                     (len(text_of_1_token), len(text_of_1_token) + 1 + len(text_of_1_token)),
                 )
 
-                tokenizer_r = self.rust_tokenizer_class.from_pretrained(
+                tokenizer_r = self.get_rust_tokenizer(
                     pretrained_name, use_fast=True, add_prefix_space=False, trim_offsets=False
                 )
                 encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
@@ -274,7 +284,7 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
                 #     (1 + len(text_of_1_token) + 1, 1 + len(text_of_1_token) + 1 + len(text_of_1_token)),
                 # )
 
-                tokenizer_r = self.rust_tokenizer_class.from_pretrained(
+                tokenizer_r = self.get_rust_tokenizer(
                     pretrained_name, use_fast=True, add_prefix_space=False, trim_offsets=True
                 )
                 encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
@@ -284,7 +294,7 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
                     (1 + len(text_of_1_token) + 1, 1 + len(text_of_1_token) + 1 + len(text_of_1_token)),
                 )
 
-                tokenizer_r = self.rust_tokenizer_class.from_pretrained(
+                tokenizer_r = self.get_rust_tokenizer(
                     pretrained_name, use_fast=True, add_prefix_space=True, trim_offsets=False
                 )
                 encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
@@ -294,7 +304,7 @@ class RobertaTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
                     (1 + len(text_of_1_token), 1 + len(text_of_1_token) + 1 + len(text_of_1_token)),
                 )
 
-                tokenizer_r = self.rust_tokenizer_class.from_pretrained(
+                tokenizer_r = self.get_rust_tokenizer(
                     pretrained_name, use_fast=True, add_prefix_space=False, trim_offsets=False
                 )
                 encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)