import torch
from mindspeed_rl.utils.pad_process import (
remove_padding_and_split_to_list,
pad_multiple,
truncate_middle_and_pad, truncate_rows,
)
from tests.test_tools.dist_test import DistributedTest
class TestPadProcess(DistributedTest):
world_size = 1
is_dist_test = False
def test_remove_padding_and_split_to_list(self):
responses = torch.tensor([[1, 2, 3, 4, 0, 0], [4, 5, 6, 1, 1, 0]], dtype=torch.int64)
eos_token_id = 77
pad_token_id = 0
expected_output = [
torch.tensor([1, 2, 3, 4], dtype=torch.int64),
torch.tensor([4, 5, 6, 1, 1], dtype=torch.int64)
]
output = remove_padding_and_split_to_list(responses, eos_token_id, pad_token_id)
assert len(output) == len(expected_output), "After processing, the length is wrong!"
for out, exp in zip(output, expected_output):
assert torch.all(out == exp), "The function of remove padding failed!"
def test_pad_multiple(self):
data_list = [
torch.tensor([1, 2, 3]),
torch.tensor([4, 5]),
torch.tensor([6])
]
pad_id = 0
multiple = 2
expected_output = torch.tensor([
[1, 2, 3, 0],
[4, 5, 0, 0],
[6, 0, 0, 0]
])
output = pad_multiple(data_list, pad_id, multiple)
assert output.shape == expected_output.shape, "After padding, the shape has been changed!"
assert torch.allclose(output, expected_output), "After padding, the result has been changed!"
def test_truncate_middle_and_pad(self):
input_tensor = torch.tensor([
[[0.1, 0.2, 0.3], [0.3, 0.4, 0.5], [0.5, 0.6, 0.7], [0.8, 0.9, 0.1]]
], dtype=torch.float32)
truncate_lengths = torch.tensor([[1, 4]], dtype=torch.int64)
pad_value = 0.0
responses = torch.tensor([2, 3, 4], dtype=torch.int64)
output = truncate_middle_and_pad(responses, input_tensor, truncate_lengths, pad_value)
expected_output = torch.tensor([
[[0.3, 0.4, 0.5], [0.5, 0.6, 0.7], [0.8, 0.9, 0.1]]
], dtype=torch.float32)
assert output.shape == expected_output.shape, "After truncating, the shape has been changed!"
assert torch.allclose(output, expected_output), "After truncating, the result has been changed!"
def test_truncate_rows(self):
tensor = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
index_tensor = torch.tensor([[2], [3], [1]])
expected_output = [torch.tensor([1, 2]), torch.tensor([5, 6, 7]), torch.tensor([9])]
result = truncate_rows(tensor, index_tensor)
for res, expected in zip(result, expected_output):
assert torch.allclose(res, expected, atol=1e-5)