Repository: OpenNMT/OpenNMT-tf
Pull request #507 (2019-10-03 10:00:43)

Fix create_variables for LM models


Patch:
diff --git a/opennmt/models/language_model.py b/opennmt/models/language_model.py
index 0be07357..7482631a 100644
--- a/opennmt/models/language_model.py
+++ b/opennmt/models/language_model.py
@@ -139,8 +139,9 @@ class LanguageModelInputter(inputters.WordEmbedder):
   input sequence.
   """
 
-  def _generate_example(self, element, training=None):
-    features = self.make_features(element, training=training)
+  def make_features(self, element=None, features=None, training=None):
+    features = super(LanguageModelInputter, self).make_features(
+        element=element, features=features, training=training)
     labels = {
         "ids_out": tf.concat([features["ids"][1:], [constants.END_OF_SENTENCE_ID]], 0),
         "length": tf.identity(features["length"])
@@ -149,6 +150,22 @@ class LanguageModelInputter(inputters.WordEmbedder):
       labels["tokens"] = tf.concat([features["tokens"][1:], [constants.END_OF_SENTENCE_TOKEN]], 0)
     return features, labels
 
+  def make_inference_dataset(self,
+                             features_file,
+                             batch_size,
+                             length_bucket_width=None,
+                             num_threads=1,
+                             prefetch_buffer_size=None):
+    dataset = self.make_dataset(features_file, training=False)
+    dataset = dataset.apply(dataset_util.inference_pipeline(
+        batch_size,
+        process_fn=lambda x: self.make_features(x, training=False)[0],
+        length_bucket_width=length_bucket_width,
+        length_fn=self.get_length,
+        num_threads=num_threads,
+        prefetch_buffer_size=prefetch_buffer_size))
+    return dataset
+
   def make_evaluation_dataset(self,
                               features_file,
                               labels_file,
@@ -160,7 +177,7 @@ class LanguageModelInputter(inputters.WordEmbedder):
     dataset = self.make_dataset(features_file, training=False)
     dataset = dataset.apply(dataset_util.inference_pipeline(
         batch_size,
-        process_fn=lambda x: self._generate_example(x, training=False),
+        process_fn=lambda x: self.make_features(x, training=False),
         num_threads=num_threads,
         prefetch_buffer_size=prefetch_buffer_size))
     return dataset
@@ -190,7 +207,7 @@ class LanguageModelInputter(inputters.WordEmbedder):
         batch_multiplier=batch_multiplier,
         length_bucket_width=length_bucket_width,
         single_pass=single_pass,
-        process_fn=lambda x: self._generate_example(x, training=True),
+        process_fn=lambda x: self.make_features(x, training=True),
         num_threads=num_threads,
         shuffle_buffer_size=shuffle_buffer_size,
         prefetch_buffer_size=prefetch_buffer_size,
diff --git a/opennmt/models/model.py b/opennmt/models/model.py
index eec965b2..a54dcdd6 100644
--- a/opennmt/models/model.py
+++ b/opennmt/models/model.py
@@ -267,7 +267,7 @@ class Model(tf.keras.layers.Layer):
             [dim or 1 for dim in spec.shape.as_list()[1:]],
             tf.constant("" if spec.dtype is tf.string else 1, dtype=spec.dtype)),
         self.examples_inputter.input_signature())
-    features = self.examples_inputter.make_features(features=features)
+    features = self.examples_inputter.make_features(features=features, training=True)
 
     # Add the batch dimension back before calling the model.
     features, labels = tf.nest.map_structure(lambda x: tf.expand_dims(x, 0), features)


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index a35f63bf..4197b4eb 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -310,6 +310,15 @@ class ModelTest(tf.test.TestCase):
     model.create_variables()
     self.assertTrue(len(model.trainable_variables) > 0)
 
+  def testCreateVariablesLanguageModel(self):
+    _, data_config = self._makeToyLMData()
+    decoder = decoders.SelfAttentionDecoder(
+        2, num_units=16, num_heads=4, ffn_inner_dim=32, num_sources=0)
+    model = models.LanguageModel(decoder, embedding_size=16)
+    model.initialize(data_config)
+    model.create_variables()
+    self.assertTrue(len(model.trainable_variables) > 0)
+
   def testInitializeWithDropoutOverride(self):
     model = models.SequenceToSequence(
         inputters.WordEmbedder(16),