@parameterized.parameters(dict(task_name='enc_dec_partial_trim_both', task_feature_lengths={'inputs': 7, 'targets': 2}, expect_trim_inputs=True, expect_trim_targets=True), dict(task_name='enc_dec_partial_trim_targets', task_feature_lengths={'inputs': None, 'targets': 2}, expect_trim_inputs=False, expect_trim_targets=True), dict(task_name='enc_dec_partial_trim_inputs', task_feature_lengths={'inputs': 7, 'targets': None}, expect_trim_inputs=True, expect_trim_targets=False), dict(task_name='enc_dec_partial_trim_neither', task_feature_lengths={'inputs': None, 'targets': None}, expect_trim_inputs=False, expect_trim_targets=False), dict(task_name='enc_dec_partial_trim_nothing', task_feature_lengths=None, expect_trim_inputs=False, expect_trim_targets=False))
def test_partial_sequence_length(self, task_name, task_feature_lengths, expect_trim_inputs, expect_trim_targets):
    x = [{'inputs': [7, 8, 5, 6, 9, 4, 3], 'targets': [3, 9]}, {'inputs': [8, 4], 'targets': [4]}, {'inputs': [5, 6, 7], 'targets': [6, 5]}]
    ds = create_default_dataset(x)
    dataset_fn = (lambda split, shuffle_files: ds)
    register_dummy_task(task_name, dataset_fn=dataset_fn)
    task = dataset_providers.get_mixture_or_task(task_name)
    output_ds = task.get_dataset(sequence_length=task_feature_lengths, shuffle=False)
    expected = [{'inputs': [7, 8, 5, 6, 9, 4, 3, 1], 'targets': [3, 9, 1]}, {'inputs': [8, 4, 1], 'targets': [4, 1]}, {'inputs': [5, 6, 7, 1], 'targets': [6, 5, 1]}]
    if expect_trim_inputs:
        expected[0]['inputs'] = [7, 8, 5, 6, 9, 4, 1]
    if expect_trim_targets:
        expected[0]['targets'] = [3, 1]
        expected[2]['targets'] = [6, 1]
    expected_dtypes = {feat: tf.int32 for feat in expected[0].keys()}
    assert_dataset(output_ds, expected, expected_dtypes=expected_dtypes)
@parameterized.parameters(dict(task_name='enc_dec_partial_trim_both', task_feature_lengths={'inputs': 7, 'targets': 2}, expect_trim_inputs=True, expect_trim_targets=True), dict(task_name='enc_dec_partial_trim_targets', task_feature_lengths={'inputs': None, 'targets': 2}, expect_trim_inputs=False, expect_trim_targets=True), dict(task_name='enc_dec_partial_trim_inputs', task_feature_lengths={'inputs': 7, 'targets': None}, expect_trim_inputs=True, expect_trim_targets=False), dict(task_name='enc_dec_partial_trim_neither', task_feature_lengths={'inputs': None, 'targets': None}, expect_trim_inputs=False, expect_trim_targets=False), dict(task_name='enc_dec_partial_trim_nothing', task_feature_lengths=None, expect_trim_inputs=False, expect_trim_targets=False))
def test_partial_sequence_length(self, task_name, task_feature_lengths, expect_trim_inputs, expect_trim_targets):
    x = [{'inputs': [7, 8, 5, 6, 9, 4, 3], 'targets': [3, 9]}, {'inputs': [8, 4], 'targets': [4]}, {'inputs': [5, 6, 7], 'targets': [6, 5]}]
    ds = create_default_dataset(x)
    dataset_fn = (lambda split, shuffle_files: ds)
    register_dummy_task(task_name, dataset_fn=dataset_fn)
    task = dataset_providers.get_mixture_or_task(task_name)
    output_ds = task.get_dataset(sequence_length=task_feature_lengths, shuffle=False)
    expected = [{'inputs': [7, 8, 5, 6, 9, 4, 3, 1], 'targets': [3, 9, 1]}, {'inputs': [8, 4, 1], 'targets': [4, 1]}, {'inputs': [5, 6, 7, 1], 'targets': [6, 5, 1]}]
    if expect_trim_inputs:
        expected[0]['inputs'] = [7, 8, 5, 6, 9, 4, 1]
    if expect_trim_targets:
        expected[0]['targets'] = [3, 1]
        expected[2]['targets'] = [6, 1]
    expected_dtypes = {feat: tf.int32 for feat in expected[0].keys()}
    assert_dataset(output_ds, expected, expected_dtypes=expected_dtypes)