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))