def test_make_positions_with_constraints_no_constraint_params_with_padding():
pad_idx = 100
test = torch.tensor([
[1, 1, 2, 3, 99, 4, 100],
[5, 6, 99, 7, 8, 100, 100],
])
expected_positions = torch.tensor([
[101, 102, 103, 104, 105, 106, 100],
[101, 102, 103, 104, 105, 100, 100],
])
calculated_positions = make_positions_with_constraints(test,
padding_idx=pad_idx)
assert torch.all(expected_positions.eq(calculated_positions))
def test_make_positions_with_constraints_no_constraint_params():
pad_idx = 0
test = torch.tensor([
[1, 1, 2, 3, 99, 4, 0],
[5, 6, 99, 7, 8, 0, 0],
])
expected_positions = torch.tensor([
[1, 2, 3, 4, 5, 6, 0],
[1, 2, 3, 4, 5, 0, 0],
])
calculated_positions = make_positions_with_constraints(test,
padding_idx=pad_idx)
assert torch.all(expected_positions.eq(calculated_positions))
def test_make_positions_with_constraints_apply_shift_and_padding():
pad_idx = 100
positional_marker_symbol_idx = 99
positional_idx_restart_offset = 10
test = torch.tensor([
[1, 1, 2, 3, 99, 4, 100],
[5, 6, 99, 7, 8, 100, 100],
])
expected_positions = torch.tensor([
[101, 102, 103, 104, 111, 112, 100],
[101, 102, 111, 112, 113, 100, 100],
])
calculated_positions = make_positions_with_constraints(
test,
padding_idx=pad_idx,
positional_marker_symbol_idx=positional_marker_symbol_idx,
positional_idx_restart_offset=positional_idx_restart_offset,
)
assert torch.all(expected_positions.eq(calculated_positions))