<reponame>shayne-longpre/BIG-bench
# 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.seqio import bigbench_bridge as bb
from bigbench.seqio import bigbench_json_paths as bb_json_paths
from bigbench.seqio import task_api
from bigbench.seqio import test_utils
from bigbench.seqio 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()

