Repository: OpenNMT/OpenNMT-tf
Pull request #908 (2021-12-14 11:35:00)

Check dataset sizes before the training


Patch:
diff --git a/opennmt/inputters/inputter.py b/opennmt/inputters/inputter.py
index 7077266e..c9fa67d6 100644
--- a/opennmt/inputters/inputter.py
+++ b/opennmt/inputters/inputter.py
@@ -424,11 +424,15 @@ class ParallelInputter(MultiInputter):
         return parallel_datasets
 
     def get_dataset_size(self, data_file):
+        common_size = None
         for inputter, data in zip(self.inputters, data_file):
             size = inputter.get_dataset_size(data)
             if size is not None:
-                return size
-        return None
+                if common_size is None:
+                    common_size = size
+                elif size != common_size:
+                    raise RuntimeError("Parallel datasets do not have the same size")
+        return common_size
 
     def input_signature(self):
         if self.combine_features:
@@ -935,6 +939,18 @@ class ExampleInputter(ParallelInputter, ExampleInputterAdapter):
             for dataset, annotation_datasets in zip(datasets, all_annotation_datasets)
         ]
 
+    def get_dataset_size(self, data_file):
+        size = super().get_dataset_size(data_file)
+        if size is not None:
+            for annotation, path in self.annotation_files.items():
+                annotation_size = tf.nest.map_structure(misc.count_lines, path)
+                if size != annotation_size:
+                    raise RuntimeError(
+                        "Annotation dataset '%s' does not have the same size as "
+                        "the examples dataset" % annotation
+                    )
+        return size
+
     def make_features(self, element=None, features=None, training=None):
         if training and self.annotation_files:
             annotations = element.copy()


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
index 0a0ee96b..48036665 100644
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -552,6 +552,51 @@ class InputterTest(tf.test.TestCase):
         self.assertAllEqual(examples[0][0]["ids"], [[0, 1, 2]])
         self.assertAllEqual(examples[0][1]["ids"], [[0]])
 
+    def testExampleInputterMismatchedDatasetSize(self):
+        vocab_file = self._makeTextFile("vocab.txt", ["a", "b", "c", "d"])
+        features_file = self._makeTextFile("features.txt", ["a", "b", "c", "d"])
+        labels_file = self._makeTextFile("labels.txt", ["a", "b", "c"])
+
+        example_inputter = inputter.ExampleInputter(
+            text_inputter.WordEmbedder(embedding_size=10),
+            text_inputter.WordEmbedder(embedding_size=10),
+        )
+        example_inputter.initialize(
+            {"source_vocabulary": vocab_file, "target_vocabulary": vocab_file}
+        )
+
+        with self.assertRaisesRegex(RuntimeError, "same size"):
+            example_inputter.make_training_dataset(
+                features_file,
+                labels_file,
+                batch_size=2,
+            )
+
+    def testExampleInputterMismatchedAnnotationsSize(self):
+        vocab_file = self._makeTextFile("vocab.txt", ["a", "b", "c", "d"])
+        features_file = self._makeTextFile("features.txt", ["a", "b", "c"])
+        labels_file = self._makeTextFile("labels.txt", ["a", "b", "c"])
+        weights_file = self._makeTextFile("weights.txt", ["1.0", "1.0"])
+
+        example_inputter = inputter.ExampleInputter(
+            text_inputter.WordEmbedder(embedding_size=10),
+            text_inputter.WordEmbedder(embedding_size=10),
+        )
+        example_inputter.initialize(
+            {
+                "source_vocabulary": vocab_file,
+                "target_vocabulary": vocab_file,
+                "example_weights": weights_file,
+            }
+        )
+
+        with self.assertRaisesRegex(RuntimeError, "same size"):
+            example_inputter.make_training_dataset(
+                features_file,
+                labels_file,
+                batch_size=2,
+            )
+
     def testWeightedDataset(self):
         vocab_file = self._makeTextFile("vocab.txt", ["the", "world", "hello", "toto"])
         data_file = self._makeTextFile("data.txt", ["hello world !"])