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)
