add wiki-scramble dataset.

PiperOrigin-RevId: 168037859

Author: nshazeer

tensor2tensor/data_generators/wiki.py

@@ -25,6 +25,8 @@
 
 import bz2file
 
+import numpy as np
+
 import six
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
@@ -130,3 +132,118 @@ def generator(self, data_dir, tmp_dir, _):
       encoded = encoder.encode(page) + [EOS]
       encoded_title = encoder.encode(title) + [EOS]
       yield {"inputs": encoded_title, "targets": encoded}
+
+
+class LanguagemodelWikiScramble(problem.Text2TextProblem):
+  """Language modeling on English wikipedia.
+
+  "targets" is a sequence of sequence_length tokens - a fragment of an article.
+  "inputs" is a copy of "targets", but with a random scramble_fraction of the
+    tokens randomly permuted.
+
+  This dataset is intended to test parallel (non-autoregressive) prediction
+  of the target sequence given the input sequence.
+  """
+
+  @property
+  def sequence_length(self):
+    raise NotImplementedError()
+
+  @property
+  def scramble_fraction(self):
+    raise NotImplementedError()
+
+  @property
+  def is_character_level(self):
+    return False
+
+  @property
+  def has_inputs(self):
+    return True
+
+  @property
+  def input_space_id(self):
+    return problem.SpaceID.EN_TOK
+
+  @property
+  def target_space_id(self):
+    return problem.SpaceID.EN_TOK
+
+  @property
+  def num_shards(self):
+    return 1000
+
+  @property
+  def vocab_name(self):
+    return "vocab.wiki"
+
+  @property
+  def use_subword_tokenizer(self):
+    return True
+
+  @property
+  def targeted_vocab_size(self):
+    return 2**13  # 8192
+
+  @property
+  def use_train_shards_for_dev(self):
+    return True
+
+  @property
+  def max_cases(self):
+    return (2 ** 30) / self.sequence_length
+
+  def scramble(self, seq):
+    seq = np.array(seq)
+    num_permute = int(self.sequence_length * self.scramble_fraction)
+    full_permutation = np.random.permutation(self.sequence_length)
+    inverse_full_permutation = np.argsort(full_permutation)
+    partial_permutation = np.random.permutation(num_permute)
+    seq = seq[full_permutation]
+    seq = np.concatenate(
+        (seq[:num_permute][partial_permutation], seq[num_permute:]))
+    seq = seq[inverse_full_permutation]
+    seq = list(seq)
+    return seq
+
+  def generator(self, data_dir, tmp_dir, _):
+    encoder = generator_utils.get_or_generate_vocab_inner(
+        data_dir, self.vocab_file, self.targeted_vocab_size,
+        lambda: page_generator(tmp_dir, max_docs=1000))
+    case_num = 0
+    for page in page_generator(tmp_dir):
+      encoded = encoder.encode(page)
+      for i in xrange(len(encoded) // self.sequence_length):
+        case_num += 1
+        if self.max_cases and case_num > self.max_cases:
+          return
+        targets = encoded[
+            i * self.sequence_length:(i + 1) * self.sequence_length]
+        inputs = self.scramble(targets)
+        yield {"inputs": inputs, "targets": targets}
+
+
+@registry.register_problem
+class LanguagemodelWikiScramble1k50(LanguagemodelWikiScramble):
+  """Sequence length 1024, 50% scrambed."""
+
+  @property
+  def sequence_length(self):
+    return 1024
+
+  @property
+  def scramble_fraction(self):
+    return 0.5
+
+
+@registry.register_problem
+class LanguagemodelWikiScramble8k50(LanguagemodelWikiScramble):
+  """Sequence length 8192, 50% scrambed."""
+
+  @property
+  def sequence_length(self):
+    return 8192
+
+  @property
+  def scramble_fraction(self):
+    return 0.5

tensor2tensor/models/attention_lm_moe.py

@@ -68,18 +68,21 @@ def model_fn_body_sharded(self, sharded_features):
     # Remove dropout if not training
     hparams = self._hparams
     dp = self._data_parallelism
-    targets = sharded_features["targets"]
-    targets = dp(tf.squeeze, targets, 2)
+    if hparams.use_inputs:
+      decoder_input = dp(tf.squeeze, sharded_features["inputs"], 2)
+      decoder_self_attention_bias = None
+    else:
+      targets = sharded_features["targets"]
+      targets = dp(tf.squeeze, targets, 2)
+      (decoder_input, decoder_self_attention_bias, pad_remover) = dp(
+          attention_lm_moe_prepare_decoder, targets, hparams)
 
     def preprocess(x):
       return dp(common_layers.layer_preprocess, x, hparams)
 
     def postprocess(x, y):
       return dp(common_layers.layer_postprocess, x, y, hparams)
 
-    (decoder_input, decoder_self_attention_bias, pad_remover) = dp(
-        attention_lm_moe_prepare_decoder, targets, hparams)
-
     x = dp(tf.nn.dropout, decoder_input,
            1.0 - hparams.layer_prepostprocess_dropout)
     extra_loss = 0.0
@@ -95,7 +98,8 @@ def _diet_expert(x):
       expert_fn = expert_utils.ffn_expert_fn(
           hparams.hidden_size, moe_hidden_sizes, hparams.hidden_size)
 
-    if hparams.attention_type == AttentionType.LOCAL_EXPERTS:
+    if (hparams.attention_type == AttentionType.LOCAL_EXPERTS
+        and not hparams.use_inputs):
       # As preprocess and postprocess are called with batch of size one (all
       # batches concatenated), we just make sure that batch_norm is not use (
       # should not either way)
@@ -162,7 +166,7 @@ def print_shape(x, suffix, debug=False):
                 attention_num_experts=hparams.attention_num_experts,
                 train=hparams.mode == ModeKeys.TRAIN,
                 batch_coordinate=batch_coordinate,
-                mask_right=True,
+                mask_right=not hparams.use_inputs,
                 split_batch=bool(hparams.attention_split_batch),
                 attention_kq_size=hparams.attention_kq_size,
                 attention_v_size=hparams.attention_v_size)
@@ -356,6 +360,9 @@ def attention_lm_moe_base():
   hparams.add_hparam("use_sepconv", int(False))
   hparams.add_hparam("diet_experts", int(False))
   hparams.add_hparam("memory_efficient_ffn", int(False))
+  # if True, we learn a non-autoregressive model from "inputs" to "targets".
+  # if False, we learn an autoregressive model to generate "targets"
+  hparams.add_hparam("use_inputs", int(False))
   return hparams
 
 
@@ -526,3 +533,17 @@ def attention_lm_moe_translation():
   hparams.moe_layers = "0,1,2,3,4,5"
   hparams.shared_embedding_and_softmax_weights = int(True)
   return hparams
+
+
+@registry.register_hparams
+def attention_lm_moe_unscramble_base():
+  """Version to use with languagemodel_wiki_scramble1k50."""
+  hparams = attention_lm_no_moe_small()
+  hparams.use_inputs = True
+  hparams.min_length_bucket = 1024
+  hparams.max_length = 1024
+  hparams.batch_size = 5000
+  hparams.layer_prepostprocess_dropout = 0.0
+  hparams.layer_preprocess_sequence = "n"
+  hparams.layer_postprocess_sequence = "da"
+  return hparams