121 lines
4.6 KiB
Python
121 lines
4.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 pytorch_wrapper.utils import pad_truncate_sequence
|
|
|
|
|
|
class TestPadTruncateSequence(unittest.TestCase):
|
|
def setUp(self):
|
|
# 3 sample sequences of different lengths
|
|
self.token_tensor_list = [
|
|
torch.tensor([1, 2, 3], dtype=torch.long),
|
|
torch.tensor([4, 5], dtype=torch.long),
|
|
torch.tensor([6], dtype=torch.long)
|
|
]
|
|
self.max_length = 255
|
|
|
|
def run_case(self, padding, padding_method, padding_value, truncation):
|
|
|
|
with self.subTest(
|
|
padding=padding,
|
|
padding_method=padding_method,
|
|
padding_value=padding_value,
|
|
truncation=truncation
|
|
):
|
|
tokens, masks = pad_truncate_sequence(
|
|
token_tensor_list=self.token_tensor_list,
|
|
padding=padding,
|
|
padding_method=padding_method,
|
|
padding_value=padding_value,
|
|
truncation=truncation,
|
|
max_length=self.max_length
|
|
)
|
|
|
|
if padding:
|
|
# Tokens and masks should be stacked tensors
|
|
self.assertIsInstance(tokens, torch.Tensor)
|
|
self.assertIsInstance(masks, torch.Tensor)
|
|
self.assertEqual(tokens.shape, masks.shape)
|
|
self.assertEqual(tokens.dim(), 2) # batch x seq_len
|
|
self.assertEqual(tokens.size(0), 3)
|
|
if padding_method == "max_length":
|
|
self.assertEqual(tokens.size(1), self.max_length)
|
|
elif padding_method == "longest":
|
|
truncated_lengths = [min(len(t), self.max_length) if truncation else len(t)
|
|
for t in self.token_tensor_list]
|
|
expected_len = max(truncated_lengths)
|
|
self.assertEqual(tokens.size(1), expected_len)
|
|
else:
|
|
# Should return lists
|
|
self.assertIsInstance(tokens, list)
|
|
self.assertIsInstance(masks, list)
|
|
self.assertEqual(len(tokens), 3)
|
|
for t, m in zip(tokens, masks):
|
|
self.assertEqual(t.size(), m.size())
|
|
expected_len = min(len(t), self.max_length) if truncation else len(t)
|
|
self.assertEqual(t.size(0), expected_len)
|
|
|
|
def test_all_cases(self):
|
|
for padding, padding_method, padding_value, truncation in product(
|
|
[True, False],
|
|
["max_length", "longest"],
|
|
[0, 255],
|
|
[True, False]
|
|
):
|
|
self.run_case(padding, padding_method, padding_value, truncation)
|
|
|
|
def test_single_tensor_input(self):
|
|
"""Test behavior when token_tensor_list has only one element."""
|
|
single_tensor_list = [torch.tensor([1, 2, 3, 4], dtype=torch.long)]
|
|
|
|
for padding, padding_method, padding_value, truncation in product(
|
|
[True, False],
|
|
["max_length", "longest"],
|
|
[0, 255],
|
|
[True, False]
|
|
):
|
|
with self.subTest(
|
|
single_input=True,
|
|
padding=padding,
|
|
padding_method=padding_method,
|
|
padding_value=padding_value,
|
|
truncation=truncation
|
|
):
|
|
tokens, masks = pad_truncate_sequence(
|
|
token_tensor_list=single_tensor_list,
|
|
padding=padding,
|
|
padding_method=padding_method,
|
|
padding_value=padding_value,
|
|
truncation=truncation,
|
|
max_length=self.max_length
|
|
)
|
|
|
|
# Always stacked even if padding is False
|
|
self.assertIsInstance(tokens, torch.Tensor)
|
|
self.assertIsInstance(masks, torch.Tensor)
|
|
self.assertEqual(tokens.shape, masks.shape)
|
|
self.assertEqual(tokens.size(0), 1)
|
|
|
|
expected_len = 4
|
|
if truncation:
|
|
expected_len = min(expected_len, self.max_length)
|
|
if padding:
|
|
if padding_method == "max_length":
|
|
expected_len = self.max_length
|
|
elif padding_method == "longest":
|
|
expected_len = expected_len # no change
|
|
|
|
self.assertEqual(tokens.size(1), expected_len)
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|