MRPC: Exclude dev data from training dataset

tensor2tensor/data_generators/mrpc.py

@@ -58,6 +58,9 @@ def dataset_splits(self):
     }, {
         "split": problem.DatasetSplit.EVAL,
         "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.TEST,
+        "shards": 1,
     }]
 
   @property
@@ -89,15 +92,18 @@ def download_file(tdir, filepath, url):
 
     return mrpc_dir
 
-  def example_generator(self, filename, dev_ids):
+  def example_generator(self, filename, dev_ids, dataset_split):
     for idx, line in enumerate(tf.gfile.Open(filename, "rb")):
       if idx == 0: continue  # skip header
       if six.PY2:
         line = unicode(line.strip(), "utf-8")
       else:
         line = line.strip().decode("utf-8")
       l, id1, id2, s1, s2 = line.split("\t")
-      if dev_ids and [id1, id2] not in dev_ids:
+      is_dev = [id1, id2] in dev_ids
+      if dataset_split == problem.DatasetSplit.TRAIN and is_dev:
+        continue
+      if dataset_split == problem.DatasetSplit.EVAL and not is_dev:
         continue
       inputs = [[s1, s2], [s2, s1]]
       for inp in inputs:
@@ -108,14 +114,17 @@ def example_generator(self, filename, dev_ids):
 
   def generate_samples(self, data_dir, tmp_dir, dataset_split):
     mrpc_dir = self._maybe_download_corpora(tmp_dir)
-    filesplit = "msr_paraphrase_train.txt"
+    if dataset_split != problem.DatasetSplit.TEST:
+      filesplit = "msr_paraphrase_train.txt"
+    else:
+      filesplit = "msr_paraphrase_test.txt"
     dev_ids = []
-    if dataset_split != problem.DatasetSplit.TRAIN:
+    if dataset_split != problem.DatasetSplit.TEST:
       for row in tf.gfile.Open(os.path.join(mrpc_dir, "dev_ids.tsv")):
         dev_ids.append(row.strip().split("\t"))
 
     filename = os.path.join(mrpc_dir, filesplit)
-    for example in self.example_generator(filename, dev_ids):
+    for example in self.example_generator(filename, dev_ids, dataset_split):
       yield example