def test_get_intra_splits_invalid_groups(self):
items_and_groups = []
# 4 groups with 2, 3, 4, 5 items each.
groups = ['A'] * 2 + ['B'] * 3 + ['C'] * 4 + ['D'] * 5
items_and_groups = list(enumerate(groups))
for split_num in range(3):
two_split_groups = [(str(i), 1.0 / 2) for i in range(2)]
res = splitter.get_splits_by_group(items_and_groups,
two_split_groups, split_num,
'intra')
self.assertEqual({'A', 'B', 'C', 'D'}, res.keys())
three_split_groups = [(str(i), 1.0 / 3) for i in range(3)]
res = splitter.get_splits_by_group(items_and_groups,
three_split_groups, split_num,
'intra')
# Now A shouldn't have a split
self.assertEqual({'B', 'C', 'D'}, res.keys())
four_split_groups = [(str(i), 1.0 / 4) for i in range(4)]
res = splitter.get_splits_by_group(items_and_groups,
four_split_groups, split_num,
'intra')
# Now A and B shouldn't have a split
self.assertEqual({'C', 'D'}, res.keys())
five_split_groups = [(str(i), 1.0 / 5) for i in range(5)]
res = splitter.get_splits_by_group(items_and_groups,
five_split_groups, split_num,
'intra')
# Now A, B and C shouldn't have a split
self.assertEqual({'D'}, res.keys())
def test_get_intra_splits_by_group_correct_splits(self):
items_and_groups = []
# 4 groups with 10, 15, 20, 25 items each.
groups = ['A'] * 10 + ['B'] * 15 + ['C'] * 20 + ['D'] * 25
items_and_groups = list(enumerate(groups))
split_probs = [('train', 0.5), ('dev', 0.25), ('test', 0.25)]
for split_num in range(3):
res = splitter.get_splits_by_group(items_and_groups, split_probs,
split_num, 'intra')
# Check we have the expected 4 groups in the result.
self.assertEqual({'A', 'B', 'C', 'D'}, res.keys())
# Check that every group has a train dev and test, and that they have the
# correct items.
i2g = dict(items_and_groups)
for group, l in [('A', 10), ('B', 15), ('C', 20), ('D', 25)]:
cur_splits = res[group]
# Check current split has train dev and test sets
self.assertEqual({'train', 'dev', 'test'}, cur_splits.keys())
# Check the current split has the correct number of items
self.assertEqual(
l, sum([len(items) for items in cur_splits.values()]))
for items in cur_splits.values():
for item in items:
self.assertEqual(group, i2g[item])
def test_get_inter_splits_by_group_correct_splits(self):
items_and_groups = [ # 4 groups with 1 or 2 items each.
(1, 'A'),
(2, 'A'),
(3, 'B'),
(4, 'B'),
(5, 'C'),
(6, 'C'),
(7, 'D'),
]
split_probs = [('train', 0.5), ('dev', 0.25), ('test', 0.25)]
for split_num in range(3):
res = splitter.get_splits_by_group(items_and_groups, split_probs,
split_num, 'inter')
# Check we have the expected 3 splits in the result.
self.assertEqual({'train', 'dev', 'test'}, res.keys())
# Check the union of all splits is all items and there are no duplicates.
self.assertCountEqual([1, 2, 3, 4, 5, 6, 7],
list(res['train']) + list(res['dev']) +
list(res['test']))
# Check items for each group appear in the same split.
i2g = dict(items_and_groups)
split_groups = { # Get list of unique groups in each split.
name: list({i2g[x]
for x in items})
for name, items in res.items()
}
self.assertCountEqual(['A', 'B', 'C', 'D'], split_groups['train'] +
split_groups['dev'] + split_groups['test'])
# Check number of groups in train/dev/test are as requested.
self.assertLen(split_groups['train'], 2)
self.assertLen(split_groups['dev'], 1)
self.assertLen(split_groups['test'], 1)