def test_get_squad_v2_twice(self):
get_squad(version=2)
with mock.patch('lineflow.datasets.squad.pickle', autospec=True) as mock_pickle:
get_squad(version=2)
mock_pickle.dump.assert_not_called()
self.assertEqual(mock_pickle.load.call_count, 1)
def test_get_squad_v2(self):
raw = get_squad(version=2)
self.assertIn('train', raw)
self.assertEqual(len(raw['train']), 130_319)
self.assertIn('dev', raw)
self.assertEqual(len(raw['dev']), 11_873)
def test_get_squad_v1(self):
raw = get_squad(version=1)
self.assertIn('train', raw)
self.assertEqual(len(raw['train']), 87_599)
self.assertIn('dev', raw)
self.assertEqual(len(raw['dev']), 10_570)