Repository: OpenNMT/OpenNMT-tf
Pull request #850 (2021-07-02 08:18:05)

Fix error when SelfAttentionDecoder is called without the source length


Patch:
diff --git a/opennmt/decoders/self_attention_decoder.py b/opennmt/decoders/self_attention_decoder.py
index 87e84156..9b4e9d21 100644
--- a/opennmt/decoders/self_attention_decoder.py
+++ b/opennmt/decoders/self_attention_decoder.py
@@ -133,13 +133,15 @@ class SelfAttentionDecoder(decoder.Decoder):
         if memory is not None:
             if not isinstance(memory, (list, tuple)):
                 memory = (memory,)
-        if memory_sequence_length is not None:
-            if not isinstance(memory_sequence_length, (list, tuple)):
-                memory_sequence_length = (memory_sequence_length,)
-            memory_mask = [
-                tf.sequence_mask(mem_length, maxlen=tf.shape(mem)[1])
-                for mem, mem_length in zip(memory, memory_sequence_length)
-            ]
+            if memory_sequence_length is not None:
+                if not isinstance(memory_sequence_length, (list, tuple)):
+                    memory_sequence_length = (memory_sequence_length,)
+                memory_mask = [
+                    tf.sequence_mask(mem_length, maxlen=tf.shape(mem)[1])
+                    for mem, mem_length in zip(memory, memory_sequence_length)
+                ]
+            else:
+                memory_mask = tuple(None for _ in memory)
 
         # Run each layer.
         new_cache = []
diff --git a/opennmt/models/catalog.py b/opennmt/models/catalog.py
index 8d80d96e..a878d4a5 100644
--- a/opennmt/models/catalog.py
+++ b/opennmt/models/catalog.py
@@ -290,65 +290,31 @@ class LstmCnnCrfTagger(sequence_tagger.SequenceTagger):
         )
 
 
-class _DefaultTransformer(transformer.Transformer):
-    def __init__(self, big=False, relative=False, shared_embeddings=False):
-        if big:
-            num_units = 1024
-            num_heads = 16
-            ffn_inner_dim = 4096
-        else:
-            num_units = 512
-            num_heads = 8
-            ffn_inner_dim = 2048
-        if relative:
-            position_encoder_class = None
-            maximum_relative_position = 20
-        else:
-            position_encoder_class = layers.SinusoidalPositionEncoder
-            maximum_relative_position = None
-        super().__init__(
-            source_inputter=inputters.WordEmbedder(embedding_size=num_units),
-            target_inputter=inputters.WordEmbedder(embedding_size=num_units),
-            num_layers=6,
-            num_units=num_units,
-            num_heads=num_heads,
-            ffn_inner_dim=ffn_inner_dim,
-            dropout=0.1,
-            attention_dropout=0.1,
-            ffn_dropout=0.1,
-            position_encoder_class=position_encoder_class,
-            share_embeddings=(
-                sequence_to_sequence.EmbeddingsSharingLevel.ALL
-                if shared_embeddings
-                else sequence_to_sequence.EmbeddingsSharingLevel.NONE
-            ),
-            maximum_relative_position=maximum_relative_position,
-        )
-
-
 @register_model_in_catalog(alias="Transformer")
-class TransformerBase(_DefaultTransformer):
+class TransformerBase(transformer.Transformer):
     """Defines a base Transformer model as described in https://arxiv.org/abs/1706.03762."""
 
 
-@register_model_in_catalog()
-class TransformerBaseSharedEmbeddings(_DefaultTransformer):
+@register_model_in_catalog
+class TransformerBaseSharedEmbeddings(transformer.Transformer):
     """Defines a base Transformer model with shared embeddings as described in
     https://arxiv.org/abs/1706.03762.
     """
 
     def __init__(self):
-        super().__init__(shared_embeddings=True)
+        super().__init__(
+            share_embeddings=sequence_to_sequence.EmbeddingsSharingLevel.ALL,
+        )
 
 
 @register_model_in_catalog(alias="TransformerRelative")
-class TransformerBaseRelative(_DefaultTransformer):
+class TransformerBaseRelative(transformer.Transformer):
     """Defines a base Transformer model using relative position representations as
     described in https://arxiv.org/abs/1803.02155.
     """
 
     def __init__(self):
-        super().__init__(relative=True)
+        super().__init__(position_encoder_class=None, maximum_relative_position=20)
 
 
 # Backward compatibility with model descriptions that directly accessed the catalog module.
@@ -357,31 +323,42 @@ TransformerRelative = TransformerBaseRelative
 
 
 @register_model_in_catalog
-class TransformerBig(_DefaultTransformer):
+class TransformerBig(transformer.Transformer):
     """Defines a big Transformer model as described in https://arxiv.org/abs/1706.03762."""
 
     def __init__(self):
-        super().__init__(big=True)
+        super().__init__(num_units=1024, num_heads=16, ffn_inner_dim=4096)
 
 
 @register_model_in_catalog
-class TransformerBigSharedEmbeddings(_DefaultTransformer):
+class TransformerBigSharedEmbeddings(transformer.Transformer):
     """Defines a big Transformer model with shared embeddings as described in
     https://arxiv.org/abs/1706.03762.
     """
 
     def __init__(self):
-        super().__init__(big=True, shared_embeddings=True)
+        super().__init__(
+            num_units=1024,
+            num_heads=16,
+            ffn_inner_dim=4096,
+            share_embeddings=sequence_to_sequence.EmbeddingsSharingLevel.ALL,
+        )
 
 
 @register_model_in_catalog
-class TransformerBigRelative(_DefaultTransformer):
+class TransformerBigRelative(transformer.Transformer):
     """Defines a big Transformer model using relative position representations as
     described in https://arxiv.org/abs/1803.02155.
     """
 
     def __init__(self):
-        super().__init__(big=True, relative=True)
+        super().__init__(
+            num_units=1024,
+            num_heads=16,
+            ffn_inner_dim=4096,
+            position_encoder_class=None,
+            maximum_relative_position=20,
+        )
 
 
 @register_model_in_catalog
diff --git a/opennmt/models/transformer.py b/opennmt/models/transformer.py
index d8d5128f..86be90f7 100644
--- a/opennmt/models/transformer.py
+++ b/opennmt/models/transformer.py
@@ -2,6 +2,7 @@
 
 import tensorflow as tf
 
+from opennmt import inputters
 from opennmt.models.sequence_to_sequence import (
     SequenceToSequence,
     EmbeddingsSharingLevel,
@@ -21,12 +22,12 @@ class Transformer(SequenceToSequence):
 
     def __init__(
         self,
-        source_inputter,
-        target_inputter,
-        num_layers,
-        num_units,
-        num_heads,
-        ffn_inner_dim,
+        source_inputter=None,
+        target_inputter=None,
+        num_layers=6,
+        num_units=512,
+        num_heads=8,
+        ffn_inner_dim=2048,
         dropout=0.1,
         attention_dropout=0.1,
         ffn_dropout=0.1,
@@ -43,10 +44,14 @@ class Transformer(SequenceToSequence):
         Args:
           source_inputter: A :class:`opennmt.inputters.Inputter` to process
             the source data. If this inputter returns parallel inputs, a multi
-            source Transformer architecture will be constructed.
+            source Transformer architecture will be constructed. Defaults to a
+            :class:`opennmt.inputters.WordEmbedder` with :obj:`num_units` as
+            embedding size.
           target_inputter: A :class:`opennmt.inputters.Inputter` to process
             the target data. Currently, only the
-            :class:`opennmt.inputters.WordEmbedder` is supported.
+            :class:`opennmt.inputters.WordEmbedder` is supported. Defaults to a
+            :class:`opennmt.inputters.WordEmbedder` with :obj:`num_units` as
+            embedding size.
           num_layers: The number of layers or a 2-tuple with the number of encoder
             layers and decoder layers.
           num_units: The number of hidden units.
@@ -76,6 +81,11 @@ class Transformer(SequenceToSequence):
             "seems better for harder-to-learn models, so it should probably be the
             default."
         """
+        if source_inputter is None:
+            source_inputter = inputters.WordEmbedder(embedding_size=num_units)
+        if target_inputter is None:
+            target_inputter = inputters.WordEmbedder(embedding_size=num_units)
+
         if isinstance(num_layers, (list, tuple)):
             num_encoder_layers, num_decoder_layers = num_layers
         else:


Test patch:
diff --git a/opennmt/tests/decoder_test.py b/opennmt/tests/decoder_test.py
index c03051ec..eb89d9d3 100644
--- a/opennmt/tests/decoder_test.py
+++ b/opennmt/tests/decoder_test.py
@@ -165,6 +165,23 @@ class DecoderTest(tf.test.TestCase):
         )
         self._testDecoder(decoder, num_sources=num_sources)
 
+    def testSelfAttentionDecoderWithoutSourceLength(self):
+        batch_size = 4
+        depth = 6
+        decoder = decoders.SelfAttentionDecoder(
+            num_layers=2,
+            num_units=depth,
+            num_heads=2,
+            ffn_inner_dim=depth * 2,
+            vocab_size=10,
+        )
+
+        memory, _, _ = _generate_source_context(batch_size, depth)
+        inputs = tf.random.uniform([batch_size, depth])
+        step = tf.constant(0)
+        initial_state = decoder.initial_state(memory)
+        decoder(inputs, step, state=initial_state)
+
 
 if __name__ == "__main__":
     tf.test.main()