Add language_model flag to train a language model by ignoring the encoder.

PiperOrigin-RevId: 182426171

Author: lmthang

nmt/model.py

@@ -179,6 +179,7 @@ def _set_params_initializer(self,
     self.tgt_vocab_size = hparams.tgt_vocab_size
     self.num_gpus = hparams.num_gpus
     self.time_major = hparams.time_major
+    self.dtype = tf.float32
 
     # extra_args: to make it flexible for adding external customizable code
     self.single_cell_fn = None
@@ -347,11 +348,14 @@ def build_graph(self, hparams, scope=None):
         bahdanau | normed_bahdanau).
     """
     utils.print_out("# creating %s graph ..." % self.mode)
-    dtype = tf.float32
 
-    with tf.variable_scope(scope or "dynamic_seq2seq", dtype=dtype):
+    with tf.variable_scope(scope or "dynamic_seq2seq", dtype=self.dtype):
       # Encoder
-      self.encoder_outputs, encoder_state = self._build_encoder(hparams)
+      if hparams.language_model:  # no encoder for language modeling
+        self.encoder_outputs = None
+        encoder_state = None
+      else:
+        self.encoder_outputs, encoder_state = self._build_encoder(hparams)
 
       ## Decoder
       logits, sample_id, final_context_state = self._build_decoder(
@@ -737,6 +741,12 @@ def _build_decoder_cell(self, hparams, encoder_outputs, encoder_state,
         base_gpu=base_gpu
     )
 
+    if hparams.language_model:
+      encoder_state = cell.zero_state(self.batch_size, self.dtype)
+    elif not hparams.pass_hidden_state:
+      raise ValueError("For non-attentional model, "
+                       "pass_hidden_state needs to be set to True")
+
     # For beam search, we need to replicate encoder infos beam_width times
     if self.mode == tf.contrib.learn.ModeKeys.INFER and hparams.beam_width > 0:
       decoder_initial_state = tf.contrib.seq2seq.tile_batch(

nmt/nmt.py

@@ -240,6 +240,9 @@ def add_arguments(parser):
                       Average the last N checkpoints for external evaluation.
                       N can be controlled by setting --num_keep_ckpts.\
                       """))
+  parser.add_argument("--language_model", type="bool", nargs="?",
+                      const=True, default=False,
+                      help="True to train a language model, ignoring encoder")
 
   # Inference
   parser.add_argument("--ckpt", type=str, default="",
@@ -369,6 +372,7 @@ def create_hparams(flags):
       override_loaded_hparams=flags.override_loaded_hparams,
       num_keep_ckpts=flags.num_keep_ckpts,
       avg_ckpts=flags.avg_ckpts,
+      language_model=flags.language_model,
       num_intra_threads=flags.num_intra_threads,
       num_inter_threads=flags.num_inter_threads,
   )
@@ -429,6 +433,16 @@ def extend_hparams(hparams):
   _add_argument(hparams, "num_decoder_residual_layers",
                 num_decoder_residual_layers)
 
+  # Language modeling
+  if hparams.language_model:
+    hparams.attention = ""
+    hparams.attention_architecture = ""
+    hparams.pass_hidden_state = False
+    hparams.share_vocab = True
+    hparams.src = hparams.tgt
+    utils.print_out("For language modeling, we turn off attention and "
+                    "pass_hidden_state; turn on share_vocab; set src to tgt.")
+
   ## Vocab
   # Get vocab file names first
   if hparams.vocab_prefix:
@@ -464,10 +478,13 @@ def extend_hparams(hparams):
   _add_argument(hparams, "src_vocab_file", src_vocab_file)
   _add_argument(hparams, "tgt_vocab_file", tgt_vocab_file)
 
-  # Pretrained Embeddings:
+  # Pretrained Embeddings
   _add_argument(hparams, "src_embed_file", "")
   _add_argument(hparams, "tgt_embed_file", "")
   if hparams.embed_prefix:
+    hparams.num_embeddings_partitions = 1
+    utils.print_out(
+        "For pretrained embeddings, set num_embeddings_partitions to 1")
     src_embed_file = hparams.embed_prefix + "." + hparams.src
     tgt_embed_file = hparams.embed_prefix + "." + hparams.tgt
 

nmt/utils/common_test_utils.py

@@ -73,6 +73,7 @@ def create_test_hparams(unit_type="lstm",
   # Misc
   standard_hparams.forget_bias = 0.0
   standard_hparams.random_seed = 3
+  language_model=False
 
   # Vocab
   standard_hparams.src_vocab_size = 5

nmt/utils/standard_hparams_utils.py

@@ -101,4 +101,7 @@ def create_standard_hparams():
       infer_batch_size=32,
       sampling_temperature=0.0,
       num_translations_per_input=1,
+
+      # Language model
+      language_model=False,
   )