def test_get_dataset_fn_train_split(self): self.mock_bigbench_json_task._ds = [{ "input": str(i), "target": [str(i + 1), str(i + 2)] } for i in range(100)] self.mock_bigbench_json_task.task_prefix = "prefix: " self.mock_bb_json_task_api.return_value = self.mock_bigbench_json_task dataset_fn = bb.get_dataset_fn(task_name="test_task", task_path="test/task/path.json") ds = dataset_fn(split="train", shuffle_files=False) expected = [ [ { # multiple targets are split into separate examples. "inputs": f"prefix: {i}", "targets": str(i + 1), "idx": np.array([idx, 0]) }, { "inputs": f"prefix: {i}", "targets": str(i + 2), "idx": np.array([idx, 1]) } ] for idx, i in enumerate(self._get_train_split_ids()) ] expected = self._flatten(expected) test_utils.assert_dataset(ds, expected)
def test_get_dataset_fn_scoring(self): # Setup mock BIG-bench task self.mock_bigbench_json_task._ds = self._flatten([[{ "input": "some input", "target": ["some output 1", "some output 2"] }, { "input": "some more input", "target": ["some more output"] }] for _ in range(100)]) self.mock_bigbench_json_task.task_prefix = "prefix: " self.mock_bb_json_task_api.return_value = self.mock_bigbench_json_task dataset_fn = bb.get_dataset_fn( task_name="test_task", task_path="test/task/path.json", bigbench_task_type=bb.BigBenchTaskType.SCORING) ds = dataset_fn(split="all", shuffle_files=False) expected = self._flatten([[{ "inputs": "prefix: some input", "targets": "some output 1", "idx": np.array([idx, 0]) }, { "inputs": "prefix: some input", "targets": "some output 2", "idx": np.array([idx, 1]) }, { "inputs": "prefix: some more input", "targets": "some more output", "idx": np.array([idx + 1, 0]) }] for idx in range(0, 200, 2)]) test_utils.assert_dataset(ds, expected)
def test_get_dataset_fn_multiple_choice_gen_targets(self): self.mock_bigbench_json_task._ds = self._flatten([[ { "input": "some input 1", "target_scores": { "ans1": 0, "ans2": 1, }, "target": ["some target 1"] }, { "input": "some input 2", "target_scores": { "ans1": 1, "ans2": 0, }, "target": ["some target 2", "some target 3"] }, ] for _ in range(100)]) self.mock_bigbench_json_task.task_prefix = "prefix: " self.mock_bb_json_task_api.return_value = self.mock_bigbench_json_task dataset_fn = bb.get_dataset_fn( task_name="test_task", task_path="test/task/path.json", bigbench_task_type=bb.BigBenchTaskType.MULTIPLE_CHOICE) ds = dataset_fn(split="all", shuffle_files=False) expected = [[{ "inputs": "prefix: some input 1", "targets": "ans1", "scores": 0.0, "idx": [idx, 0], "answers": ["some target 1"] }, { "inputs": "prefix: some input 1", "targets": "ans2", "scores": 1.0, "idx": [idx, 1], "answers": ["some target 1"] }, { "inputs": "prefix: some input 2", "targets": "ans1", "scores": 1.0, "idx": [idx + 1, 0], "answers": ["some target 2", "some target 3"] }, { "inputs": "prefix: some input 2", "targets": "ans2", "scores": 0.0, "idx": [idx + 1, 1], "answers": ["some target 2", "some target 3"] }] for idx in range(0, 200, 2)] expected = self._flatten(expected) test_utils.assert_dataset(ds, expected)
def test_get_dataset_fn_validation_split(self): self.mock_bigbench_json_task._ds = [{ "input": str(i), "target": [str(i + 1), str(i + 2)] } for i in range(100)] self.mock_bigbench_json_task.task_prefix = "prefix: " self.mock_bb_json_task_api.return_value = self.mock_bigbench_json_task dataset_fn = bb.get_dataset_fn(task_name="test_task", task_path="test/task/path.json") ds = dataset_fn(split="validation", shuffle_files=False) expected = [{ # multiple targets are not split into separate examples. "inputs": f"prefix: {i}", "targets": str(i + 1), "answers": [str(i + 1), str(i + 2)], } for i in self._get_validation_split_ids()] test_utils.assert_dataset(ds, expected)
def test_get_dataset_fn_all_split(self): self.mock_bigbench_json_task._ds = [{ "input": str(i), "target": [str(i + 1)] } for i in range(100)] self.mock_bigbench_json_task.task_prefix = "prefix: " self.mock_bb_json_task_api.return_value = self.mock_bigbench_json_task dataset_fn = bb.get_dataset_fn(task_name="test_task", task_path="test/task/path.json") ds = dataset_fn(split="all", shuffle_files=False) expected = [{ "inputs": f"prefix: {i}", "targets": str(i + 1), "answers": [str(i + 1)], } for i in range(100)] test_utils.assert_dataset(ds, expected)
def test_get_dataset_fn_few_shot(self, mock_nshot_maker): self.mock_bigbench_json_task._ds = [{ "input": "some input", "target": ["some output"] } for _ in range(100)] self.mock_bigbench_json_task.few_shot_example_separator = "," self.mock_bigbench_json_task.task_prefix = "prefix: " self.mock_bigbench_json_task.rng = None self.mock_bb_json_task_api.return_value = self.mock_bigbench_json_task mock_nshot_maker.return_value = [{ "input": "fewshot input", "target": ["fewshot output"] } for _ in range(100)] dataset_fn = bb.get_dataset_fn(task_name="test_task", task_path="test/task/path.json", num_shots=1) ds = dataset_fn(split="all", shuffle_files=False) expected = [{ "inputs": "prefix: fewshot input", "targets": "fewshot output", "answers": ["fewshot output"] } for _ in range(100)] test_utils.assert_dataset(ds, expected)
def test_get_dataset_fn_shuffle_max_examples(self): # Setup mock BIG-bench task self.mock_bigbench_json_task._ds = [{ "input": str(i), "target": [str(i + 1)] } for i in range(100)] self.mock_bigbench_json_task.task_prefix = "" self.mock_bb_json_task_api.return_value = self.mock_bigbench_json_task dataset_fn = bb.get_dataset_fn(task_name="test_task", task_path="test/task/path.json", max_examples=50) ds_all = dataset_fn(split="all", shuffle_files=True, seed=23) ds_train = dataset_fn(split="train", shuffle_files=True, seed=23) ds_val = dataset_fn(split="validation", shuffle_files=True, seed=23) expected_all = [{ "inputs": str(i), "targets": str(i + 1), "answers": [str(i + 1)], } for i in self._get_shuffled_sampled_all_split_ids()] expected_train = [{ "inputs": str(i), "targets": str(i + 1), "idx": np.array([idx, 0]) } for idx, i in enumerate(self._get_shuffled_sampled_train_split_ids()) ] expected_val = [{ "inputs": str(i), "targets": str(i + 1), "answers": [str(i + 1)], } for i in self._get_shuffled_sampled_validation_split_ids()] test_utils.assert_dataset(ds_all, expected_all) test_utils.assert_dataset(ds_train, expected_train) test_utils.assert_dataset(ds_val, expected_val)