internal merge from github

PiperOrigin-RevId: 160210398

Author: Lukasz Kaiser

tensor2tensor/data_generators/wmt.py

@@ -29,7 +29,12 @@
 import tensorflow as tf
 
 
-def character_generator(source_path, target_path, eos=None):
+# End-of-sentence marker (should correspond to the position of EOS in the
+# RESERVED_TOKENS list in text_encoder.py)
+EOS = 1
+
+
+def character_generator(source_path, target_path, character_vocab, eos=None):
   """Generator for sequence-to-sequence tasks that just uses characters.
 
   This generator assumes the files at source_path and target_path have
@@ -40,6 +45,7 @@ def character_generator(source_path, target_path, eos=None):
   Args:
     source_path: path to the file with source sentences.
     target_path: path to the file with target sentences.
+    character_vocab: a TextEncoder to encode the characters.
     eos: integer to append at the end of each sequence (default: None).
 
   Yields:
@@ -51,8 +57,8 @@ def character_generator(source_path, target_path, eos=None):
     with tf.gfile.GFile(target_path, mode="r") as target_file:
       source, target = source_file.readline(), target_file.readline()
       while source and target:
-        source_ints = [ord(c) for c in source.strip()] + eos_list
-        target_ints = [ord(c) for c in target.strip()] + eos_list
+        source_ints = character_vocab.encode(source.strip()) + eos_list
+        target_ints = character_vocab.encode(target.strip()) + eos_list
         yield {"inputs": source_ints, "targets": target_ints}
         source, target = source_file.readline(), target_file.readline()
 
@@ -226,14 +232,16 @@ def ende_wordpiece_token_generator(tmp_dir, train, vocab_size):
   tag = "train" if train else "dev"
   data_path = _compile_data(tmp_dir, datasets, "wmt_ende_tok_%s" % tag)
   return token_generator(data_path + ".lang1", data_path + ".lang2",
-                         symbolizer_vocab, 1)
+                         symbolizer_vocab, EOS)
 
 
 def ende_character_generator(tmp_dir, train):
+  character_vocab = text_encoder.ByteTextEncoder()
   datasets = _ENDE_TRAIN_DATASETS if train else _ENDE_TEST_DATASETS
   tag = "train" if train else "dev"
   data_path = _compile_data(tmp_dir, datasets, "wmt_ende_chr_%s" % tag)
-  return character_generator(data_path + ".lang1", data_path + ".lang2", 1)
+  return character_generator(data_path + ".lang1", data_path + ".lang2",
+                             character_vocab, EOS)
 
 
 def enfr_wordpiece_token_generator(tmp_dir, train, vocab_size):
@@ -244,22 +252,25 @@ def enfr_wordpiece_token_generator(tmp_dir, train, vocab_size):
   tag = "train" if train else "dev"
   data_path = _compile_data(tmp_dir, datasets, "wmt_enfr_tok_%s" % tag)
   return token_generator(data_path + ".lang1", data_path + ".lang2",
-                         symbolizer_vocab, 1)
+                         symbolizer_vocab, EOS)
 
 
 def enfr_character_generator(tmp_dir, train):
   """Instance of character generator for the WMT en->fr task."""
+  character_vocab = text_encoder.ByteTextEncoder()
   datasets = _ENFR_TRAIN_DATASETS if train else _ENFR_TEST_DATASETS
   tag = "train" if train else "dev"
   data_path = _compile_data(tmp_dir, datasets, "wmt_enfr_chr_%s" % tag)
-  return character_generator(data_path + ".lang1", data_path + ".lang2", 1)
+  return character_generator(data_path + ".lang1", data_path + ".lang2",
+                             character_vocab, EOS)
 
 
 def parsing_character_generator(tmp_dir, train):
+  character_vocab = text_encoder.ByteTextEncoder()
   filename = "parsing_%s" % ("train" if train else "dev")
   text_filepath = os.path.join(tmp_dir, filename + ".text")
   tags_filepath = os.path.join(tmp_dir, filename + ".tags")
-  return character_generator(text_filepath, tags_filepath, 1)
+  return character_generator(text_filepath, tags_filepath, character_vocab, EOS)
 
 
 def parsing_token_generator(tmp_dir, train, vocab_size):
@@ -268,4 +279,4 @@ def parsing_token_generator(tmp_dir, train, vocab_size):
   filename = "parsing_%s" % ("train" if train else "dev")
   text_filepath = os.path.join(tmp_dir, filename + ".text")
   tags_filepath = os.path.join(tmp_dir, filename + ".tags")
-  return token_generator(text_filepath, tags_filepath, symbolizer_vocab, 1)
+  return token_generator(text_filepath, tags_filepath, symbolizer_vocab, EOS)

tensor2tensor/data_generators/wmt_test.py

@@ -25,6 +25,7 @@
 # Dependency imports
 
 import six
+from tensor2tensor.data_generators import text_encoder
 from tensor2tensor.data_generators import wmt
 
 import tensorflow as tf
@@ -36,31 +37,48 @@ def testCharacterGenerator(self):
     # Generate a trivial source and target file.
     tmp_dir = self.get_temp_dir()
     (_, tmp_file_path) = tempfile.mkstemp(dir=tmp_dir)
+    if six.PY2:
+      enc_f = lambda s: s
+    else:
+      enc_f = lambda s: s.encode("utf-8")
     with io.open(tmp_file_path + ".src", "wb") as src_file:
-      src_file.write("source1\n")
-      src_file.write("source2\n")
+      src_file.write(enc_f("source1\n"))
+      src_file.write(enc_f("source2\n"))
     with io.open(tmp_file_path + ".tgt", "wb") as tgt_file:
-      tgt_file.write("target1\n")
-      tgt_file.write("target2\n")
+      tgt_file.write(enc_f("target1\n"))
+      tgt_file.write(enc_f("target2\n"))
 
     # Call character generator on the generated files.
     results_src, results_tgt = [], []
-    for dictionary in wmt.character_generator(tmp_file_path + ".src",
-                                              tmp_file_path + ".tgt"):
+    character_vocab = text_encoder.ByteTextEncoder()
+    for dictionary in wmt.character_generator(
+        tmp_file_path + ".src", tmp_file_path + ".tgt", character_vocab):
       self.assertEqual(sorted(list(dictionary)), ["inputs", "targets"])
       results_src.append(dictionary["inputs"])
       results_tgt.append(dictionary["targets"])
 
     # Check that the results match the files.
+    # First check that the results match the encoded original strings;
+    # this is a comparison of integer arrays.
     self.assertEqual(len(results_src), 2)
-    self.assertEqual("".join([six.int2byte(i)
-                              for i in results_src[0]]), "source1")
-    self.assertEqual("".join([six.int2byte(i)
-                              for i in results_src[1]]), "source2")
-    self.assertEqual("".join([six.int2byte(i)
-                              for i in results_tgt[0]]), "target1")
-    self.assertEqual("".join([six.int2byte(i)
-                              for i in results_tgt[1]]), "target2")
+    self.assertEqual(results_src[0],
+                     character_vocab.encode("source1"))
+    self.assertEqual(results_src[1],
+                     character_vocab.encode("source2"))
+    self.assertEqual(results_tgt[0],
+                     character_vocab.encode("target1"))
+    self.assertEqual(results_tgt[1],
+                     character_vocab.encode("target2"))
+    # Then decode the results and compare with the original strings;
+    # this is a comparison of strings
+    self.assertEqual(character_vocab.decode(results_src[0]),
+                     "source1")
+    self.assertEqual(character_vocab.decode(results_src[1]),
+                     "source2")
+    self.assertEqual(character_vocab.decode(results_tgt[0]),
+                     "target1")
+    self.assertEqual(character_vocab.decode(results_tgt[1]),
+                     "target2")
 
     # Clean up.
     os.remove(tmp_file_path + ".src")