98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import torch
|
|
from itertools import product
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
|
|
from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
|
|
|
|
|
|
class TestHfTokenizerEncode(unittest.TestCase):
|
|
def setUp(self):
|
|
self.node = HfTokenizerEncode()
|
|
self.test_sentence_1d = "Hello, world"
|
|
self.test_sentence_2d = [
|
|
"a b c d e f g h i j k l m n o p q r s t u v w x y z",
|
|
"a b c d e f g h i j k l"]
|
|
self.models = {
|
|
"bert-base-uncased": "BertTokenizer",
|
|
"roberta-base": "RobertaTokenizer",
|
|
"distilbert-base-uncased": "DistilBertTokenizer",
|
|
"albert-base-v2": "AlbertTokenizer",
|
|
"google/electra-small-discriminator": "ElectraTokenizer",
|
|
}
|
|
|
|
def test_tokenizer_variants(self):
|
|
for model_name, tokenizer_name in self.models.items():
|
|
for is_1d, padding, truncation, pad_method in product(
|
|
[True, False], [True, False], [True, False], ["longest", "max_length"]
|
|
):
|
|
|
|
if not is_1d and not padding:
|
|
continue # Cannot batch ragged output into tensors
|
|
|
|
with self.subTest(
|
|
tokenizer_name=tokenizer_name,
|
|
is_1d=is_1d,
|
|
padding=padding,
|
|
truncation=truncation,
|
|
pad_method=pad_method
|
|
):
|
|
sentence = self.test_sentence_1d if is_1d else self.test_sentence_2d
|
|
|
|
if truncation:
|
|
max_length = 10
|
|
else:
|
|
max_length = 256
|
|
|
|
encoder = self.node.f(
|
|
model_name=model_name,
|
|
padding=padding,
|
|
padding_method=pad_method,
|
|
truncation=truncation,
|
|
max_length=max_length
|
|
)[0]
|
|
output = encoder(sentence=sentence)
|
|
|
|
input_ids = output["input_ids"]
|
|
mask = output["attention_mask"]
|
|
|
|
# shape agreement
|
|
self.assertEqual(input_ids.shape, mask.shape)
|
|
|
|
# correct batch size
|
|
self.assertEqual(input_ids.shape[0], 1 if is_1d else len(sentence))
|
|
|
|
# correct max length behavior
|
|
if not is_1d:
|
|
if truncation:
|
|
self.assertEqual(input_ids.shape[1], 10)
|
|
elif padding and pad_method == "longest":
|
|
self.assertEqual(input_ids[0].size(0), input_ids[1].size(0))
|
|
elif padding and pad_method == "max_length":
|
|
self.assertGreaterEqual(input_ids.shape[1], 256)
|
|
|
|
|
|
def test_padding_false_1d(self):
|
|
for model_name in self.models:
|
|
with self.subTest(model=model_name):
|
|
encoder = self.node.f(
|
|
model_name=model_name,
|
|
padding=False,
|
|
padding_method="longest",
|
|
truncation=False,
|
|
max_length=256
|
|
)[0]
|
|
output = encoder(sentence=self.test_sentence_1d)
|
|
input_ids = output["input_ids"]
|
|
self.assertEqual(input_ids.shape[0], 1)
|
|
self.assertGreater(torch.count_nonzero(input_ids), 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|