Repository: OpenNMT/OpenNMT-tf
Pull request #971 (2022-09-16 10:20:44)

Add parameter to automatically share embeddings when possible


Patch:
diff --git a/opennmt/inputters/inputter.py b/opennmt/inputters/inputter.py
index e0d26876..c7d0487d 100644
--- a/opennmt/inputters/inputter.py
+++ b/opennmt/inputters/inputter.py
@@ -355,21 +355,10 @@ class ParallelInputter(MultiInputter):
           combine_features: Combine each inputter features in a single dict or
             return them separately. This is typically ``True`` for multi source
             inputs but ``False`` for features/labels parallel data.
-
-        Raises:
-          ValueError: if :obj:`share_parameters` is set but the child inputters are
-            not of the same type.
         """
         super().__init__(inputters, reducer=reducer)
         self.combine_features = combine_features
         self.share_parameters = share_parameters
-        if self.share_parameters:
-            leaves = self.get_leaf_inputters()
-            for inputter in leaves[1:]:
-                if type(inputter) is not type(leaves[0]):  # noqa: E721
-                    raise ValueError(
-                        "Each inputter must be of the same type for parameter sharing"
-                    )
 
     def _structure(self):
         """Returns the nested structure that represents this parallel inputter."""
@@ -545,6 +534,11 @@ class ParallelInputter(MultiInputter):
             # all attributes with parameters to the other inputters.
             leaves = self.get_leaf_inputters()
             first, others = leaves[0], leaves[1:]
+            for inputter in others:
+                if type(inputter) is not type(first):  # noqa: E721
+                    raise ValueError(
+                        "Each inputter must be of the same type for parameter sharing"
+                    )
             first.build(input_shape)
             for name, attr in first.__dict__.copy().items():
                 if isinstance(attr, tf.Variable) or (
diff --git a/opennmt/models/sequence_to_sequence.py b/opennmt/models/sequence_to_sequence.py
index c13eb279..543376c8 100644
--- a/opennmt/models/sequence_to_sequence.py
+++ b/opennmt/models/sequence_to_sequence.py
@@ -21,12 +21,14 @@ class EmbeddingsSharingLevel(object):
      * ``SOURCE_TARGET_INPUT``: share source and target word embeddings
      * ``TARGET``: share target word embeddings and softmax weights
      * ``ALL``: share words embeddings and softmax weights
+     * ``AUTO``: automatically share embeddings when using the same vocabulary file.
     """
 
     NONE = 0
     SOURCE_TARGET_INPUT = 1
     TARGET = 2
     ALL = 3
+    AUTO = 4
 
     @staticmethod
     def share_input_embeddings(level):
@@ -69,30 +71,12 @@ class SequenceToSequence(model.SequenceGenerator):
 
         Raises:
           TypeError: if :obj:`target_inputter` is not a
-            :class:`opennmt.inputters.WordEmbedder` (same for
-            :obj:`source_inputter` when embeddings sharing is enabled).
+            :class:`opennmt.inputters.WordEmbedder`.
         """
         if not isinstance(target_inputter, inputters.WordEmbedder):
             raise TypeError("Target inputter must be a WordEmbedder")
-        if EmbeddingsSharingLevel.share_input_embeddings(share_embeddings):
-            if isinstance(source_inputter, inputters.ParallelInputter):
-                source_inputters = source_inputter.inputters
-            else:
-                source_inputters = [source_inputter]
-            for inputter in source_inputters:
-                if not isinstance(inputter, inputters.WordEmbedder):
-                    raise TypeError(
-                        "Sharing embeddings requires all inputters to be a "
-                        "WordEmbedder"
-                    )
 
-        examples_inputter = SequenceToSequenceInputter(
-            source_inputter,
-            target_inputter,
-            share_parameters=EmbeddingsSharingLevel.share_input_embeddings(
-                share_embeddings
-            ),
-        )
+        examples_inputter = SequenceToSequenceInputter(source_inputter, target_inputter)
         super().__init__(examples_inputter)
         self.encoder = encoder
         self.decoder = decoder
@@ -149,6 +133,30 @@ class SequenceToSequence(model.SequenceGenerator):
             )
             self.labels_inputter.set_noise(noiser, in_place=False)
 
+        if self.share_embeddings != EmbeddingsSharingLevel.NONE:
+            all_inputters = self.examples_inputter.get_leaf_inputters()
+
+            if self.share_embeddings == EmbeddingsSharingLevel.AUTO:
+                if all(
+                    isinstance(inputter, inputters.WordEmbedder)
+                    and inputter.vocabulary_file == all_inputters[0].vocabulary_file
+                    for inputter in all_inputters
+                ):
+                    self.share_embeddings = EmbeddingsSharingLevel.ALL
+                else:
+                    self.share_embeddings = EmbeddingsSharingLevel.TARGET
+
+            if EmbeddingsSharingLevel.share_input_embeddings(self.share_embeddings):
+                if not all(
+                    isinstance(inputter, inputters.WordEmbedder)
+                    for inputter in all_inputters
+                ):
+                    raise TypeError(
+                        "Sharing embeddings requires all inputters to be a WordEmbedder"
+                    )
+
+                self.examples_inputter.share_parameters = True
+
     def build(self, input_shape):
         super().build(input_shape)
         if EmbeddingsSharingLevel.share_target_embeddings(self.share_embeddings):


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index b40f9fbb..e0cc1a36 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -229,22 +229,38 @@ class ModelTest(tf.test.TestCase):
             params=params,
         )
 
-    def testSequenceToSequenceWithSharedEmbedding(self):
+    @parameterized.expand(
+        [
+            (models.EmbeddingsSharingLevel.ALL, True, True, True),
+            (models.EmbeddingsSharingLevel.AUTO, True, True, True),
+            (models.EmbeddingsSharingLevel.AUTO, False, False, True),
+        ]
+    )
+    def testSequenceToSequenceWithSharedEmbedding(
+        self, share_embeddings, reuse_vocab, input_is_shared, target_is_shared
+    ):
         model = models.SequenceToSequence(
             inputters.WordEmbedder(16),
             inputters.WordEmbedder(16),
             encoders.SelfAttentionEncoder(2, 16, 4, 32),
             decoders.SelfAttentionDecoder(2, 16, 4, 32),
-            share_embeddings=models.EmbeddingsSharingLevel.ALL,
+            share_embeddings=share_embeddings,
         )
         _, _, data_config = self._makeToyEnDeData()
-        data_config["target_vocabulary"] = data_config["source_vocabulary"]
+        if reuse_vocab:
+            data_config["target_vocabulary"] = data_config["source_vocabulary"]
         model.initialize(data_config)
-        self.assertTrue(model.decoder.initialized)
-        model.build(None)
+        model.create_variables()
+
+        self.assertEqual(
+            model.features_inputter.embedding.ref()
+            == model.labels_inputter.embedding.ref(),
+            input_is_shared,
+        )
         self.assertEqual(
-            model.labels_inputter.embedding.ref(),
-            model.decoder.output_layer.weight.ref(),
+            model.labels_inputter.embedding.ref()
+            == model.decoder.output_layer.kernel.ref(),
+            target_is_shared,
         )
 
     @parameterized.expand(