def test_pooling_batcher(self):
data = [
RowData({"text": "something"}, {
"a": i,
"b": 10 + i,
"c": 20 + i
}) for i in range(10)
]
batcher = PoolingBatcher(train_batch_size=3, pool_num_batches=2)
batches = list(
batcher.batchify(data, sort_key=lambda x: x.numberized["a"]))
self.assertEqual(len(batches), 4)
a_vals = {a for raw_batch, batch in batches for a in batch["a"]}
self.assertSetEqual(a_vals, set(range(10)))
for raw_batch, batch in batches[:2]:
self.assertEqual([{
"text": "something"
}] * len(raw_batch), list(raw_batch))
self.assertGreater(batch["a"][0], batch["a"][-1])
for a in batch["a"]:
self.assertLess(a, 6)
for _, batch in batches[2:]:
for a in batch["a"]:
self.assertGreaterEqual(a, 6)
def _yield_and_reset(self, row):
packed_tokens = list(self.remainder["tokens"])
packed_segments = list(self.remainder["segment_labels"])
self.remainder: Dict[str, List[int]] = {
"tokens": [],
"segment_labels": []
}
return RowData(
row,
self._format_output_row(packed_tokens, packed_segments,
len(packed_tokens)),
)
def _yield_and_reset(self):
packed_tokens = list(self.remainder["tokens"])
packed_segments = list(self.remainder["segment_labels"])
self.remainder: Dict[str, List[int]] = {
"tokens": [],
"segment_labels": []
}
return RowData(
{}, # packed LM data doesn't respect data cardinality
self._format_output_row(packed_tokens, packed_segments,
len(packed_tokens)),
)
def test_batcher(self):
data = [
RowData({"text": "something"}, {
"a": i,
"b": 10 + i,
"c": 20 + i
}) for i in range(10)
]
batcher = Batcher(train_batch_size=3)
batches = list(batcher.batchify(data))
self.assertEqual(len(batches), 4)
self.assertEqual(len(batches[0].raw_data), 3)
self.assertEqual("something", batches[1].raw_data[0]["text"])
self.assertEqual(batches[1].numberized["a"], [3, 4, 5])
self.assertEqual(batches[3].numberized["b"], [19])