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:
Yih-Dar
2025-03-28 15:09:35 +01:00
committed by GitHub
parent 3af425d4c6
commit 1fcaad6df9
92 changed files with 1301 additions and 884 deletions

View File

@@ -13,12 +13,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import unittest
from functools import lru_cache
from transformers import CohereTokenizerFast
from transformers.testing_utils import require_jinja, require_tokenizers, require_torch_multi_gpu
from ...test_tokenization_common import TokenizerTesterMixin
from ...test_tokenization_common import TokenizerTesterMixin, use_cache_if_possible
@require_tokenizers
@@ -37,14 +39,21 @@ class CohereTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
"pad_token": "<PAD>",
}
def setUp(self):
super().setUp()
@classmethod
def setUpClass(cls):
super().setUpClass()
tokenizer = CohereTokenizerFast.from_pretrained("hf-internal-testing/tiny-random-CohereForCausalLM")
tokenizer.save_pretrained(self.tmpdirname)
tokenizer.save_pretrained(cls.tmpdirname)
def get_rust_tokenizer(self, **kwargs):
kwargs.update(self.special_tokens_map)
return CohereTokenizerFast.from_pretrained(self.tmpdirname, **kwargs)
@classmethod
@use_cache_if_possible
@lru_cache(maxsize=64)
def get_rust_tokenizer(cls, pretrained_name=None, **kwargs):
_kwargs = copy.deepcopy(cls.special_tokens_map)
_kwargs.update(kwargs)
kwargs = _kwargs
pretrained_name = pretrained_name or cls.tmpdirname
return CohereTokenizerFast.from_pretrained(pretrained_name, **kwargs)
# This gives CPU OOM on a single-gpu runner (~60G RAM). On multi-gpu runner, it has ~180G RAM which is enough.
@require_torch_multi_gpu
@@ -80,7 +89,7 @@ class CohereTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
def test_padding(self, max_length=10):
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_r = self.get_rust_tokenizer(pretrained_name, **kwargs)
# tokenizer_r.pad_token = None # Hotfixing padding = None
# Simple input
s = "This is a simple input"