# Copyright 2021 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for task_api.""" import itertools from absl.testing import absltest from absl.testing import parameterized from bigbench.bbseqio import bigbench_bridge as bb from bigbench.bbseqio import bigbench_json_paths as bb_json_paths from bigbench.bbseqio import task_api from bigbench.bbseqio import test_utils from bigbench.bbseqio import vocabs import seqio mock = absltest.mock GEN = bb.BigBenchTaskType.GENERATIVE SCR = bb.BigBenchTaskType.SCORING MUL = bb.BigBenchTaskType.MULTIPLE_CHOICE MIX = bb.BigBenchTaskType.MIX TEST_FEWSHOT_SETTINGS = (0, 1, 2, 3) TEST_MAX_EXAMPLES_PER_TASK = (None, 1024) TEST_VOCABS = vocabs.ALL_VOCABS TEST_PREPROCESSING_SETTINGS = { 'strip_inputs': True, 'strip_targets': True, 'add_inputs_eos': False, 'add_targets_eos': False, } class TaskApiTest(absltest.TestCase): def test_get_seqio_name(self): mock_vocab = task_api.SeqIOVocabulary( 'mock', 'A Mock Vocab', mock.create_autospec(seqio.Vocabulary)) self.assertEqual( task_api.get_seqio_name('task-name', bb.BigBenchTaskType.GENERATIVE, mock_vocab, 2), 'bigbench:task_name.gen.mock_vocab.2_shot.all_examples') self.assertEqual( task_api.get_seqio_name('task-name', bb.BigBenchTaskType.MULTIPLE_CHOICE, mock_vocab, 2), 'bigbench:task_name.mul.mock_vocab.2_shot.all_examples') self.assertEqual( task_api.get_seqio_name('task-name', bb.BigBenchTaskType.GENERATIVE, mock_vocab, 2, 'subtask'), 'bigbench:task_name.gen.mock_vocab.2_shot.all_examples.subtask') class TaskRegistrationTest(test_utils.BigbenchBridgeTest): def get_seqio_name_and_rate(self, bb_task_name, bigbench_task_type, bigbench_subtask_name, shots, vocab, max_examples): if bigbench_subtask_name: task_name = f'{bb_task_name}:{bigbench_subtask_name}' num_examples = bb_json_paths.get_num_examples(task_name) else: num_examples = bb_json_paths.get_num_examples(bb_task_name) # The following assumes that values in GLOBAL_MAX_EXAMPLES_PER_TASK > 100 max_e = num_examples if max_examples else None rate = num_examples seqio_name = task_api.get_seqio_name( bb_task_name, bigbench_task_type, vocab=vocab, num_shots=shots, bigbench_subtask_name=bigbench_subtask_name, max_examples=max_e) return (seqio_name, rate) @parameterized.named_parameters([ (f'{s}_shot_{v.name}_vocab_{m}_max', s, v, m) # pylint:disable=g-complex-comprehension for s, v, m in itertools.product(TEST_FEWSHOT_SETTINGS, TEST_VOCABS.values(), TEST_MAX_EXAMPLES_PER_TASK) ]) def test_simple_gen_scr(self, shots, vocab, max_examples): test_json_paths = {'test_simple_gen_scr': '/foo1'} test_json_metadata = { 'test_simple_gen_scr': { 'subtask_names': [], 'num_generate_text': 100, 'num_multiple_choice': 0, }, } test_json_subtask_names = {} self.mock_bb_json_paths_utils(test_json_paths, test_json_metadata, test_json_subtask_names) # Register BIG-bench tasks and mixtures, similar to tasks.py. task_names = bb_json_paths.get_task_names() task_rates = {} for task_name in task_names: _, rates = task_api.register_composite_task( task_name=task_name, shots=shots, vocab=vocab, max_examples=max_examples, **TEST_PREPROCESSING_SETTINGS) task_rates.update(rates) task_rates = sorted(task_rates.items()) # Verify registered tasks and mixtures. registered_task_names = sorted( [n for n in seqio.TaskRegistry.names() if 'bigbench:test_' in n]) registered_mixture_names = sorted( [n for n in seqio.MixtureRegistry.names() if 'bigbench:test_' in n]) expected_task_rates = sorted([ self.get_seqio_name_and_rate('test_simple_gen_scr', GEN, None, shots, vocab, max_examples), self.get_seqio_name_and_rate('test_simple_gen_scr', SCR, None, shots, vocab, max_examples), ]) expected_mixture_rates = [ self.get_seqio_name_and_rate('test_simple_gen_scr', MIX, None, shots, vocab, max_examples), ] expected_task_names = sorted( [name_and_rate[0] for name_and_rate in expected_task_rates]) expected_mixture_names = sorted( [name_and_rate[0] for name_and_rate in expected_mixture_rates]) self.assertListEqual(expected_task_names, registered_task_names) self.assertListEqual(expected_mixture_names, registered_mixture_names) # Verify the rates. self.assertListEqual(expected_task_rates, task_rates) @parameterized.named_parameters([ (f'{s}_shot_{v.name}_vocab_{m}_max', s, v, m) # pylint:disable=g-complex-comprehension for s, v, m in itertools.product(TEST_FEWSHOT_SETTINGS, TEST_VOCABS.values(), TEST_MAX_EXAMPLES_PER_TASK) ]) def test_simple_mul(self, shots, vocab, max_examples): test_json_paths = {'test_simple_mul': '/foo2'} test_json_metadata = { 'test_simple_mul': { 'subtask_names': [], 'num_generate_text': 0, 'num_multiple_choice': 100, }, } test_json_subtask_names = {} self.mock_bb_json_paths_utils(test_json_paths, test_json_metadata, test_json_subtask_names) # Register BIG-bench tasks and mixtures, similar to tasks.py. task_names = bb_json_paths.get_task_names() task_rates = {} for task_name in task_names: _, rates = task_api.register_composite_task( task_name=task_name, shots=shots, vocab=vocab, max_examples=max_examples, **TEST_PREPROCESSING_SETTINGS) task_rates.update(rates) task_rates = sorted(task_rates.items()) # Verify registered tasks and mixtures. registered_task_names = sorted( [n for n in seqio.TaskRegistry.names() if 'bigbench:test_' in n]) registered_mixture_names = sorted( [n for n in seqio.MixtureRegistry.names() if 'bigbench:test_' in n]) expected_task_rates = sorted([ self.get_seqio_name_and_rate('test_simple_mul', MUL, None, shots, vocab, max_examples), ]) expected_task_names = sorted( [name_and_rate[0] for name_and_rate in expected_task_rates]) expected_mixture_names = [] self.assertListEqual(expected_task_names, registered_task_names) self.assertListEqual(expected_mixture_names, registered_mixture_names) # Verify the rates. self.assertListEqual(expected_task_rates, task_rates) @parameterized.named_parameters([ (f'{s}_shot_{v.name}_vocab_{m}_max', s, v, m) # pylint:disable=g-complex-comprehension for s, v, m in itertools.product(TEST_FEWSHOT_SETTINGS, TEST_VOCABS.values(), TEST_MAX_EXAMPLES_PER_TASK) ]) def test_simple_combination(self, shots, vocab, max_examples): test_json_paths = { 'test_simple_combination': '/foo3', } test_json_metadata = { 'test_simple_combination': { 'subtask_names': [], 'num_generate_text': 100, 'num_multiple_choice': 100, }, } test_json_subtask_names = {} self.mock_bb_json_paths_utils(test_json_paths, test_json_metadata, test_json_subtask_names) # Register BIG-bench tasks and mixtures, similar to tasks.py. task_names = bb_json_paths.get_task_names() task_rates = {} for task_name in task_names: _, rates = task_api.register_composite_task( task_name=task_name, shots=shots, vocab=vocab, max_examples=max_examples, **TEST_PREPROCESSING_SETTINGS) task_rates.update(rates) task_rates = sorted(task_rates.items()) # Verify registered tasks and mixtures. registered_task_names = sorted( [n for n in seqio.TaskRegistry.names() if 'bigbench:test_' in n]) registered_mixture_names = sorted( [n for n in seqio.MixtureRegistry.names() if 'bigbench:test_' in n]) expected_task_rates = sorted([ self.get_seqio_name_and_rate('test_simple_combination', GEN, None, shots, vocab, max_examples), self.get_seqio_name_and_rate('test_simple_combination', MUL, None, shots, vocab, max_examples), self.get_seqio_name_and_rate('test_simple_combination', SCR, None, shots, vocab, max_examples), ]) expected_mixture_rates = [ self.get_seqio_name_and_rate('test_simple_combination', MIX, None, shots, vocab, max_examples), ] expected_task_names = sorted( [name_and_rate[0] for name_and_rate in expected_task_rates]) expected_mixture_names = sorted( [name_and_rate[0] for name_and_rate in expected_mixture_rates]) self.assertListEqual(expected_task_names, registered_task_names) self.assertListEqual(expected_mixture_names, registered_mixture_names) # Verify the rates. self.assertListEqual(expected_task_rates, task_rates) @parameterized.named_parameters([ (f'{s}_shot_{v.name}_vocab_{m}_max', s, v, m) # pylint:disable=g-complex-comprehension for s, v, m in itertools.product(TEST_FEWSHOT_SETTINGS, TEST_VOCABS.values(), TEST_MAX_EXAMPLES_PER_TASK) ]) def test_composite_gen_scr(self, shots, vocab, max_examples): test_json_paths = {'test_composite_gen_scr': '/foo4'} test_json_metadata = { 'test_composite_gen_scr': { 'subtask_names': [ 'test_composite_gen_scr:st_1', 'test_composite_gen_scr:st_2', ], 'num_generate_text': 100, 'num_multiple_choice': 0, }, 'test_composite_gen_scr:st_1': { 'subtask_names': [], 'num_generate_text': 50, 'num_multiple_choice': 0, }, 'test_composite_gen_scr:st_2': { 'subtask_names': [], 'num_generate_text': 50, 'num_multiple_choice': 0, }, } test_json_subtask_names = { 'test_composite_gen_scr': [ 'test_composite_gen_scr:st_1', 'test_composite_gen_scr:st_2', ], } self.mock_bb_json_paths_utils(test_json_paths, test_json_metadata, test_json_subtask_names) # Register BIG-bench tasks and mixtures, similar to tasks.py. task_names = bb_json_paths.get_task_names() task_rates = {} for task_name in task_names: _, rates = task_api.register_composite_task( task_name=task_name, shots=shots, vocab=vocab, max_examples=max_examples, **TEST_PREPROCESSING_SETTINGS) task_rates.update(rates) task_rates = sorted(task_rates.items()) # Verify registered tasks and mixtures. registered_task_names = sorted( [n for n in seqio.TaskRegistry.names() if 'bigbench:test_' in n]) registered_mixture_names = sorted( [n for n in seqio.MixtureRegistry.names() if 'bigbench:test_' in n]) expected_task_rates = sorted([ self.get_seqio_name_and_rate('test_composite_gen_scr', GEN, 'st_1', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_gen_scr', SCR, 'st_1', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_gen_scr', GEN, 'st_2', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_gen_scr', SCR, 'st_2', shots, vocab, max_examples), ]) expected_mixture_rates = [ self.get_seqio_name_and_rate('test_composite_gen_scr', MIX, 'st_1', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_gen_scr', MIX, 'st_2', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_gen_scr', MIX, None, shots, vocab, max_examples), ] expected_task_names = sorted( [name_and_rate[0] for name_and_rate in expected_task_rates]) expected_mixture_names = sorted( [name_and_rate[0] for name_and_rate in expected_mixture_rates]) self.assertListEqual(expected_task_names, registered_task_names) self.assertListEqual(expected_mixture_names, registered_mixture_names) # Verify the rates. self.assertListEqual(expected_task_rates, task_rates) @parameterized.named_parameters([ (f'{s}_shot_{v.name}_vocab_{m}_max', s, v, m) # pylint:disable=g-complex-comprehension for s, v, m in itertools.product(TEST_FEWSHOT_SETTINGS, TEST_VOCABS.values(), TEST_MAX_EXAMPLES_PER_TASK) ]) def test_composite_mul(self, shots, vocab, max_examples): test_json_paths = {'test_composite_mul': '/foo5'} test_json_metadata = { 'test_composite_mul': { 'subtask_names': [ 'test_composite_mul:st_1', 'test_composite_mul:st_2', ], 'num_generate_text': 0, 'num_multiple_choice': 100, }, 'test_composite_mul:st_1': { 'subtask_names': [], 'num_generate_text': 0, 'num_multiple_choice': 50, }, 'test_composite_mul:st_2': { 'subtask_names': [], 'num_generate_text': 0, 'num_multiple_choice': 50, }, } test_json_subtask_names = { 'test_composite_mul': [ 'test_composite_mul:st_1', 'test_composite_mul:st_2', ], } self.mock_bb_json_paths_utils(test_json_paths, test_json_metadata, test_json_subtask_names) # Register BIG-bench tasks and mixtures, similar to tasks.py. task_names = bb_json_paths.get_task_names() task_rates = {} for task_name in task_names: _, rates = task_api.register_composite_task( task_name=task_name, shots=shots, vocab=vocab, max_examples=max_examples, **TEST_PREPROCESSING_SETTINGS) task_rates.update(rates) task_rates = sorted(task_rates.items()) # Verify registered tasks and mixtures. registered_task_names = sorted( [n for n in seqio.TaskRegistry.names() if 'bigbench:test_' in n]) registered_mixture_names = sorted( [n for n in seqio.MixtureRegistry.names() if 'bigbench:test_' in n]) expected_task_rates = sorted([ self.get_seqio_name_and_rate('test_composite_mul', MUL, 'st_1', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_mul', MUL, 'st_2', shots, vocab, max_examples), ]) expected_mixture_rates = [ self.get_seqio_name_and_rate('test_composite_mul', MIX, None, shots, vocab, max_examples), ] expected_task_names = sorted( [name_and_rate[0] for name_and_rate in expected_task_rates]) expected_mixture_names = sorted( [name_and_rate[0] for name_and_rate in expected_mixture_rates]) self.assertListEqual(expected_task_names, registered_task_names) self.assertListEqual(expected_mixture_names, registered_mixture_names) # Verify the rates. self.assertListEqual(expected_task_rates, task_rates) @parameterized.named_parameters([ (f'{s}_shot_{v.name}_vocab_{m}_max', s, v, m) # pylint:disable=g-complex-comprehension for s, v, m in itertools.product(TEST_FEWSHOT_SETTINGS, TEST_VOCABS.values(), TEST_MAX_EXAMPLES_PER_TASK) ]) def test_composite_combination(self, shots, vocab, max_examples): test_json_paths = {'test_composite_combination': '/foo6'} test_json_metadata = { 'test_composite_combination': { 'subtask_names': [ 'test_composite_combination:gen_scr_st_1', 'test_composite_combination:mul_st_1', ], 'num_generate_text': 50, 'num_multiple_choice': 50, }, 'test_composite_combination:gen_scr_st_1': { 'subtask_names': [], 'num_generate_text': 50, 'num_multiple_choice': 0, }, 'test_composite_combination:mul_st_1': { 'subtask_names': [], 'num_generate_text': 0, 'num_multiple_choice': 50, }, } test_json_subtask_names = { 'test_composite_combination': [ 'test_composite_combination:gen_scr_st_1', 'test_composite_combination:mul_st_1', ], } self.mock_bb_json_paths_utils(test_json_paths, test_json_metadata, test_json_subtask_names) # Register BIG-bench tasks and mixtures, similar to tasks.py. task_names = bb_json_paths.get_task_names() task_rates = {} for task_name in task_names: _, rates = task_api.register_composite_task( task_name=task_name, shots=shots, vocab=vocab, max_examples=max_examples, **TEST_PREPROCESSING_SETTINGS) task_rates.update(rates) task_rates = sorted(task_rates.items()) # Verify registered tasks and mixtures. registered_task_names = sorted( [n for n in seqio.TaskRegistry.names() if 'bigbench:test_' in n]) registered_mixture_names = sorted( [n for n in seqio.MixtureRegistry.names() if 'bigbench:test_' in n]) expected_task_rates = sorted([ self.get_seqio_name_and_rate('test_composite_combination', GEN, 'gen_scr_st_1', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_combination', SCR, 'gen_scr_st_1', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_combination', MUL, 'mul_st_1', shots, vocab, max_examples), ]) expected_mixture_rates = [ self.get_seqio_name_and_rate('test_composite_combination', MIX, 'gen_scr_st_1', shots, vocab, max_examples), self.get_seqio_name_and_rate('test_composite_combination', MIX, None, shots, vocab, max_examples), ] expected_task_names = sorted( [name_and_rate[0] for name_and_rate in expected_task_rates]) expected_mixture_names = sorted( [name_and_rate[0] for name_and_rate in expected_mixture_rates]) self.assertListEqual(expected_task_names, registered_task_names) self.assertListEqual(expected_mixture_names, registered_mixture_names) # Verify the rates. self.assertListEqual(expected_task_rates, task_rates) if __name__ == '__main__': absltest.main()