Add TransformerEncoder and TransformerDecoder models

PiperOrigin-RevId: 164785525

Author: Ryan Sepassi

tensor2tensor/bin/t2t-datagen

@@ -82,6 +82,16 @@ _SUPPORTED_PROBLEM_GENERATORS = {
     "algorithmic_algebra_inverse": (
         lambda: algorithmic_math.algebra_inverse(26, 0, 2, 100000),
         lambda: algorithmic_math.algebra_inverse(26, 3, 3, 10000)),
+    "ice_parsing_tokens": (
+        lambda: wmt.tabbed_parsing_token_generator(
+            FLAGS.data_dir, FLAGS.tmp_dir, True, "ice", 2**13, 2**8),
+        lambda: wmt.tabbed_parsing_token_generator(
+            FLAGS.data_dir, FLAGS.tmp_dir, False, "ice", 2**13, 2**8)),
+    "ice_parsing_characters": (
+        lambda: wmt.tabbed_parsing_character_generator(
+            FLAGS.data_dir, FLAGS.tmp_dir, True),
+        lambda: wmt.tabbed_parsing_character_generator(
+            FLAGS.data_dir, FLAGS.tmp_dir, False)),
     "wmt_parsing_tokens_8k": (
         lambda: wmt.parsing_token_generator(
             FLAGS.data_dir, FLAGS.tmp_dir, True, 2**13),

tensor2tensor/bin/t2t-trainer

@@ -31,7 +31,6 @@
 from tensor2tensor.data_generators import wiki
 from tensor2tensor.data_generators import wmt
 from tensor2tensor.data_generators import wsj_parsing
-from tensor2tensor.data_generators import ice_parsing
 
 
 # Problem modules that require optional dependencies

tensor2tensor/data_generators/all_problems.py

@@ -462,6 +462,39 @@ def wsj_parsing_tokens(model_hparams, prefix, wrong_source_vocab_size,
   return p
 
 
+def ice_parsing_tokens(model_hparams, wrong_source_vocab_size):
+  """Icelandic to parse tree translation benchmark.
+
+  Args:
+    model_hparams: a tf.contrib.training.HParams
+    wrong_source_vocab_size: a number used in the filename indicating the
+      approximate vocabulary size.  This is not to be confused with the actual
+      vocabulary size.
+
+  Returns:
+    A tf.contrib.training.HParams object.
+  """
+  p = default_problem_hparams()
+  # This vocab file must be present within the data directory.
+  source_vocab_filename = os.path.join(
+      model_hparams.data_dir, "ice_source.vocab.%d" % wrong_source_vocab_size)
+  target_vocab_filename = os.path.join(model_hparams.data_dir,
+                                       "ice_target.vocab.256")
+  source_subtokenizer = text_encoder.SubwordTextEncoder(source_vocab_filename)
+  target_subtokenizer = text_encoder.SubwordTextEncoder(target_vocab_filename)
+  p.input_modality = {
+      "inputs": (registry.Modalities.SYMBOL, source_subtokenizer.vocab_size)
+  }
+  p.target_modality = (registry.Modalities.SYMBOL, 256)
+  p.vocabulary = {
+      "inputs": source_subtokenizer,
+      "targets": target_subtokenizer,
+  }
+  p.input_space_id = 18  # Icelandic tokens
+  p.target_space_id = 19  # Icelandic parse tokens
+  return p
+
+
 def img2img_imagenet(unused_model_hparams):
   """Image 2 Image for imagenet dataset."""
   p = default_problem_hparams()
@@ -511,6 +544,10 @@ def image_celeba(unused_model_hparams):
         lm1b_32k,
     "wiki_32k":
         wiki_32k,
+    "ice_parsing_characters":
+        wmt_parsing_characters,
+    "ice_parsing_tokens":
+        lambda p: ice_parsing_tokens(p, 2**13),
     "wmt_parsing_tokens_8k":
         lambda p: wmt_parsing_tokens(p, 2**13),
     "wsj_parsing_tokens_16k":

tensor2tensor/data_generators/generator_utils.py

@@ -648,6 +648,28 @@ def target_space_id(self):
     return problem.SpaceID.CS_CHR
 
 
+def tabbed_parsing_token_generator(data_dir, tmp_dir, train, prefix,
+                                   source_vocab_size, target_vocab_size):
+  """Generate source and target data from a single file."""
+  source_vocab = generator_utils.get_or_generate_tabbed_vocab(
+      data_dir, tmp_dir, "parsing_train.pairs", 0,
+      prefix + "_source.vocab.%d" % source_vocab_size, source_vocab_size)
+  target_vocab = generator_utils.get_or_generate_tabbed_vocab(
+      data_dir, tmp_dir, "parsing_train.pairs", 1,
+      prefix + "_target.vocab.%d" % target_vocab_size, target_vocab_size)
+  filename = "parsing_%s" % ("train" if train else "dev")
+  pair_filepath = os.path.join(tmp_dir, filename + ".pairs")
+  return tabbed_generator(pair_filepath, source_vocab, target_vocab, EOS)
+
+
+def tabbed_parsing_character_generator(tmp_dir, train):
+  """Generate source and target data from a single file."""
+  character_vocab = text_encoder.ByteTextEncoder()
+  filename = "parsing_%s" % ("train" if train else "dev")
+  pair_filepath = os.path.join(tmp_dir, filename + ".pairs")
+  return tabbed_generator(pair_filepath, character_vocab, character_vocab, EOS)
+
+
 def parsing_token_generator(data_dir, tmp_dir, train, vocab_size):
   symbolizer_vocab = generator_utils.get_or_generate_vocab(
       data_dir, tmp_dir, "vocab.endefr.%d" % vocab_size, vocab_size)

tensor2tensor/data_generators/ice_parsing.py

@@ -55,22 +55,66 @@ def model_fn_body(self, features):
     (decoder_input, decoder_self_attention_bias) = transformer_prepare_decoder(
         targets, hparams)
 
-    encoder_input = tf.nn.dropout(
-        encoder_input, 1.0 - hparams.layer_prepostprocess_dropout)
-    decoder_input = tf.nn.dropout(
-        decoder_input, 1.0 - hparams.layer_prepostprocess_dropout)
-    encoder_output = transformer_encoder(
-        encoder_input, encoder_self_attention_bias, hparams)
+    encoder_input = tf.nn.dropout(encoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+    decoder_input = tf.nn.dropout(decoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+    encoder_output = transformer_encoder(encoder_input,
+                                         encoder_self_attention_bias, hparams)
 
     decoder_output = transformer_decoder(
-        decoder_input, encoder_output,
-        decoder_self_attention_bias,
+        decoder_input, encoder_output, decoder_self_attention_bias,
         encoder_decoder_attention_bias, hparams)
     decoder_output = tf.expand_dims(decoder_output, 2)
 
     return decoder_output
 
 
+@registry.register_model
+class TransformerEncoder(t2t_model.T2TModel):
+  """Transformer, encoder only."""
+
+  def model_fn_body(self, features):
+    hparams = self._hparams
+    inputs = features["inputs"]
+    target_space = features["target_space_id"]
+
+    inputs = common_layers.flatten4d3d(inputs)
+
+    (encoder_input, encoder_self_attention_bias,
+     _) = (transformer_prepare_encoder(inputs, target_space, hparams))
+
+    encoder_input = tf.nn.dropout(encoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+    encoder_output = transformer_encoder(encoder_input,
+                                         encoder_self_attention_bias, hparams)
+
+    return encoder_output
+
+
+@registry.register_model
+class TransformerDecoder(t2t_model.T2TModel):
+  """Transformer, decoder only."""
+
+  def model_fn_body(self, features):
+    hparams = self._hparams
+    targets = features["targets"]
+
+    targets = common_layers.flatten4d3d(targets)
+
+    (decoder_input, decoder_self_attention_bias) = transformer_prepare_decoder(
+        targets, hparams)
+
+    decoder_input = tf.nn.dropout(decoder_input,
+                                  1.0 - hparams.layer_prepostprocess_dropout)
+
+    decoder_output = transformer_decoder(
+        decoder_input, None, decoder_self_attention_bias, None, hparams)
+    decoder_output = tf.expand_dims(decoder_output, 2)
+
+    return decoder_output
+
+
 def transformer_prepare_encoder(inputs, target_space, hparams):
   """Prepare one shard of the model for the encoder.
 
@@ -150,14 +194,11 @@ def transformer_encoder(encoder_input,
       with tf.variable_scope("layer_%d" % layer):
         with tf.variable_scope("self_attention"):
           y = common_attention.multihead_attention(
-              common_layers.layer_preprocess(x, hparams),
-              None,
-              encoder_self_attention_bias,
+              common_layers.layer_preprocess(
+                  x, hparams), None, encoder_self_attention_bias,
               hparams.attention_key_channels or hparams.hidden_size,
               hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout)
+              hparams.hidden_size, hparams.num_heads, hparams.attention_dropout)
           x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
@@ -196,26 +237,23 @@ def transformer_decoder(decoder_input,
       with tf.variable_scope("layer_%d" % layer):
         with tf.variable_scope("self_attention"):
           y = common_attention.multihead_attention(
-              common_layers.layer_preprocess(x, hparams),
-              None,
-              decoder_self_attention_bias,
-              hparams.attention_key_channels or hparams.hidden_size,
-              hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout)
-          x = common_layers.layer_postprocess(x, y, hparams)
-        with tf.variable_scope("encdec_attention"):
-          y = common_attention.multihead_attention(
-              common_layers.layer_preprocess(x, hparams),
-              encoder_output,
-              encoder_decoder_attention_bias,
+              common_layers.layer_preprocess(
+                  x, hparams), None, decoder_self_attention_bias,
               hparams.attention_key_channels or hparams.hidden_size,
               hparams.attention_value_channels or hparams.hidden_size,
-              hparams.hidden_size,
-              hparams.num_heads,
-              hparams.attention_dropout)
+              hparams.hidden_size, hparams.num_heads, hparams.attention_dropout)
           x = common_layers.layer_postprocess(x, y, hparams)
+        if encoder_output is not None:
+          assert encoder_decoder_attention_bias is not None
+          with tf.variable_scope("encdec_attention"):
+            y = common_attention.multihead_attention(
+                common_layers.layer_preprocess(
+                    x, hparams), encoder_output, encoder_decoder_attention_bias,
+                hparams.attention_key_channels or hparams.hidden_size,
+                hparams.attention_value_channels or hparams.hidden_size,
+                hparams.hidden_size, hparams.num_heads,
+                hparams.attention_dropout)
+            x = common_layers.layer_postprocess(x, y, hparams)
         with tf.variable_scope("ffn"):
           y = transformer_ffn_layer(
               common_layers.layer_preprocess(x, hparams), hparams)
@@ -393,7 +431,7 @@ def transformer_parsing_big():
 
 @registry.register_hparams
 def transformer_parsing_ice():
-  """Hparams for parsing and tagging Icelandic text."""
+  """Hparams for parsing Icelandic text."""
   hparams = transformer_base_single_gpu()
   hparams.batch_size = 4096
   hparams.shared_embedding_and_softmax_weights = int(False)

tensor2tensor/data_generators/problem_hparams.py

@@ -259,19 +259,14 @@ def _interactive_input_fn(hparams):
   vocabulary = p_hparams.vocabulary["inputs" if has_input else "targets"]
   # This should be longer than the longest input.
   const_array_size = 10000
-  # Import readline if available for command line editing and recall
-  try:
-    import readline
-  except ImportError:
-    pass
   while True:
     prompt = ("INTERACTIVE MODE  num_samples=%d  decode_length=%d  \n"
               "  it=<input_type>     ('text' or 'image' or 'label')\n"
               "  pr=<problem_num>    (set the problem number)\n"
               "  in=<input_problem>  (set the input problem number)\n"
               "  ou=<output_problem> (set the output problem number)\n"
               "  ns=<num_samples>    (changes number of samples)\n"
-              "  dl=<decode_length>  (changes decode length)\n"
+              "  dl=<decode_length>  (changes decode legnth)\n"
               "  <%s>                (decode)\n"
               "  q                   (quit)\n"
               ">" % (num_samples, decode_length, "source_string"