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:
@@ -17,6 +17,7 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from functools import lru_cache
|
||||
from typing import List
|
||||
|
||||
from transformers.models.esm.tokenization_esm import VOCAB_FILES_NAMES, EsmTokenizer
|
||||
@@ -24,24 +25,32 @@ from transformers.testing_utils import require_tokenizers
|
||||
from transformers.tokenization_utils import PreTrainedTokenizer
|
||||
from transformers.tokenization_utils_base import PreTrainedTokenizerBase
|
||||
|
||||
from ...test_tokenization_common import use_cache_if_possible
|
||||
|
||||
|
||||
@require_tokenizers
|
||||
class ESMTokenizationTest(unittest.TestCase):
|
||||
tokenizer_class = EsmTokenizer
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.tmpdirname = tempfile.mkdtemp()
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
|
||||
cls.tmpdirname = tempfile.mkdtemp()
|
||||
vocab_tokens: List[str] = ["<cls>", "<pad>", "<eos>", "<unk>", "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", "O", ".", "-", "<null_1>", "<mask>"] # fmt: skip
|
||||
self.vocab_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
|
||||
with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
|
||||
cls.vocab_file = os.path.join(cls.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
|
||||
with open(cls.vocab_file, "w", encoding="utf-8") as vocab_writer:
|
||||
vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
|
||||
|
||||
def get_tokenizers(self, **kwargs) -> List[PreTrainedTokenizerBase]:
|
||||
return [self.get_tokenizer(**kwargs)]
|
||||
def get_tokenizers(cls, **kwargs) -> List[PreTrainedTokenizerBase]:
|
||||
return [cls.get_tokenizer(**kwargs)]
|
||||
|
||||
def get_tokenizer(self, **kwargs) -> PreTrainedTokenizer:
|
||||
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) -> PreTrainedTokenizer:
|
||||
pretrained_name = pretrained_name or cls.tmpdirname
|
||||
return cls.tokenizer_class.from_pretrained(pretrained_name, **kwargs)
|
||||
|
||||
def test_tokenizer_single_example(self):
|
||||
tokenizer = self.tokenizer_class(self.vocab_file)
|
||||
|
||||
Reference in New Issue
Block a user