| 607 | self.assertEqual(self.sample_examples, ds.examples) |
| 608 | |
| 609 | def test_write_roundtrip_indexed(self): |
| 610 | tempdir = self.create_tempdir() |
| 611 | output_base = os.path.join( |
| 612 | tempdir.full_path, 'test_dataset.indexed.lit.jsonl' |
| 613 | ) |
| 614 | lit_dataset.write_examples( |
| 615 | self.indexed_dataset.indexed_examples, output_base |
| 616 | ) |
| 617 | lit_dataset.write_spec(self.data_spec, output_base + '.spec') |
| 618 | |
| 619 | # Read back and compare contents |
| 620 | ds = lit_dataset.load_lit_format(output_base) |
| 621 | self.assertIsInstance(ds, lit_dataset.IndexedDataset) |
| 622 | self.assertEqual(self.data_spec, ds.spec()) |
| 623 | self.assertEqual(self.indexed_dataset.indexed_examples, ds.indexed_examples) |
| 624 | for original, loaded in zip(self.sample_examples, ds.examples): |
| 625 | for key in self.indexed_dataset.spec().keys(): |
| 626 | self.assertEqual(original[key], loaded[key]) |
| 627 | |
| 628 | |
| 629 | if __name__ == '__main__': |