Internal change

PiperOrigin-RevId: 440245619

official/nlp/data/classifier_data_lib.py

@@ -187,6 +187,8 @@ class AxProcessor(DataProcessor):
 
   def _create_examples_tfds(self, dataset, set_type):
     """Creates examples for the training/dev/test sets."""
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -218,6 +220,8 @@ class ColaProcessor(DefaultGLUEDataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/cola", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -312,6 +316,8 @@ class MnliProcessor(DataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/mnli", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -343,6 +349,8 @@ class MrpcProcessor(DefaultGLUEDataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/mrpc", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -453,6 +461,8 @@ class QnliProcessor(DefaultGLUEDataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/qnli", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -484,6 +494,8 @@ class QqpProcessor(DefaultGLUEDataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/qqp", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -517,6 +529,8 @@ class RteProcessor(DefaultGLUEDataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/rte", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -548,6 +562,8 @@ class SstProcessor(DefaultGLUEDataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/sst2", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -574,6 +590,8 @@ class StsBProcessor(DefaultGLUEDataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/stsb", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)
@@ -742,6 +760,8 @@ class WnliProcessor(DefaultGLUEDataProcessor):
     """Creates examples for the training/dev/test sets."""
     dataset = tfds.load(
         "glue/wnli", split=set_type, try_gcs=True).as_numpy_iterator()
+    dataset = list(dataset)
+    dataset.sort(key=lambda x: x["idx"])
     examples = []
     for i, example in enumerate(dataset):
       guid = "%s-%s" % (set_type, i)