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:
@@ -15,10 +15,11 @@
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from functools import lru_cache
|
||||
|
||||
from transformers.models.bertweet.tokenization_bertweet import VOCAB_FILES_NAMES, BertweetTokenizer
|
||||
|
||||
from ...test_tokenization_common import TokenizerTesterMixin
|
||||
from ...test_tokenization_common import TokenizerTesterMixin, use_cache_if_possible
|
||||
|
||||
|
||||
class BertweetTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
|
||||
@@ -26,26 +27,31 @@ class BertweetTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
|
||||
tokenizer_class = BertweetTokenizer
|
||||
test_rust_tokenizer = False
|
||||
|
||||
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 = ["I", "m", "V@@", "R@@", "r", "e@@"]
|
||||
vocab_tokens = dict(zip(vocab, range(len(vocab))))
|
||||
merges = ["#version: 0.2", "a m</w>"]
|
||||
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:
|
||||
for token in vocab_tokens:
|
||||
fp.write(f"{token} {vocab_tokens[token]}\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 BertweetTokenizer.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 BertweetTokenizer.from_pretrained(pretrained_name, **kwargs)
|
||||
|
||||
def get_input_output_texts(self, tokenizer):
|
||||
input_text = "I am VinAI Research"
|
||||
|
||||
Reference in New Issue
Block a user