1# Copyright 2020 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15""" 16Testing Mask op in DE 17""" 18import numpy as np 19 20import mindspore.dataset as ds 21import mindspore.dataset.text as text 22 23 24def compare(in1, in2, length, out1, out2): 25 data = ds.NumpySlicesDataset({"s1": [in1], "s2": [in2]}) 26 data = data.map(operations=text.TruncateSequencePair(length), input_columns=["s1", "s2"]) 27 data = data.map(input_columns=["s1", "s2"], operations=text.TruncateSequencePair(length)) 28 for d in data.create_dict_iterator(num_epochs=1, output_numpy=True): 29 np.testing.assert_array_equal(out1, d["s1"]) 30 np.testing.assert_array_equal(out2, d["s2"]) 31 32 33def test_callable(): 34 op = text.TruncateSequencePair(3) 35 data = [["1", "2", "3"], ["4", "5"]] 36 result_text = op(*data) 37 column1, column2 = op(["1", "2", "3"], ["4", "5"]) 38 assert np.array_equal(result_text[0], ['1', '2']) 39 assert np.array_equal(result_text[1], ['4']) 40 assert np.array_equal(column1, ['1', '2']) 41 assert np.array_equal(column2, ['4']) 42 43 44def test_basics(): 45 compare(in1=[1, 2, 3], in2=[4, 5], length=4, out1=[1, 2], out2=[4, 5]) 46 compare(in1=[1, 2], in2=[4, 5], length=4, out1=[1, 2], out2=[4, 5]) 47 compare(in1=[1], in2=[4], length=4, out1=[1], out2=[4]) 48 compare(in1=[1, 2, 3, 4], in2=[5], length=4, out1=[1, 2, 3], out2=[5]) 49 compare(in1=[1, 2, 3, 4], in2=[5, 6, 7, 8], length=4, out1=[1, 2], out2=[5, 6]) 50 51 52def test_basics_odd(): 53 compare(in1=[1, 2, 3], in2=[4, 5], length=3, out1=[1, 2], out2=[4]) 54 compare(in1=[1, 2], in2=[4, 5], length=3, out1=[1, 2], out2=[4]) 55 compare(in1=[1], in2=[4], length=5, out1=[1], out2=[4]) 56 compare(in1=[1, 2, 3, 4], in2=[5], length=3, out1=[1, 2], out2=[5]) 57 compare(in1=[1, 2, 3, 4], in2=[5, 6, 7, 8], length=3, out1=[1, 2], out2=[5]) 58 59 60def test_basics_str(): 61 compare(in1=[b"1", b"2", b"3"], in2=[4, 5], length=4, out1=[b"1", b"2"], out2=[4, 5]) 62 compare(in1=[b"1", b"2"], in2=[b"4", b"5"], length=4, out1=[b"1", b"2"], out2=[b"4", b"5"]) 63 compare(in1=[b"1"], in2=[4], length=4, out1=[b"1"], out2=[4]) 64 compare(in1=[b"1", b"2", b"3", b"4"], in2=[b"5"], length=4, out1=[b"1", b"2", b"3"], out2=[b"5"]) 65 compare(in1=[b"1", b"2", b"3", b"4"], in2=[5, 6, 7, 8], length=4, out1=[b"1", b"2"], out2=[5, 6]) 66 67 68def test_exceptions(): 69 compare(in1=[1, 2, 3, 4], in2=[5, 6, 7, 8], length=1, out1=[1], out2=[]) 70 71 72if __name__ == "__main__": 73 test_callable() 74 test_basics() 75 test_basics_odd() 76 test_basics_str() 77 test_exceptions() 78