imdb dataset: S3 version (issue #737).

PiperOrigin-RevId: 257574363

Author: pierrot0

tensorflow_datasets/text/imdb.py

@@ -66,6 +66,9 @@ def __init__(self, text_encoder_config=None, **kwargs):
     super(IMDBReviewsConfig, self).__init__(
         version=tfds.core.Version(
             "0.1.0", experiments={tfds.core.Experiment.S3: False}),
+        supported_versions=[
+            tfds.core.Version("1.0.0"),
+        ],
         **kwargs)
     self.text_encoder_config = (
         text_encoder_config or tfds.features.text.TextEncoderConfig())
@@ -118,8 +121,8 @@ def _info(self):
     )
 
   def _vocab_text_gen(self, archive):
-    for ex in self._generate_examples(archive,
-                                      os.path.join("aclImdb", "train")):
+    for ex in self._generate_examples(
+        archive, os.path.join("aclImdb", "train"), keys=False):
       yield ex["text"]
 
   def _split_generators(self, dl_manager):
@@ -149,7 +152,7 @@ def _split_generators(self, dl_manager):
                         "labeled": False}),
     ]
 
-  def _generate_examples(self, archive, directory, labeled=True):
+  def _generate_examples(self, archive, directory, labeled=True, keys=True):
     """Generate IMDB examples."""
     # For labeled examples, extract the label from the path.
     reg_path = "(?P<label>neg|pos)" if labeled else "unsup"
@@ -161,7 +164,11 @@ def _generate_examples(self, archive, directory, labeled=True):
         continue
       text = imdb_f.read().strip()
       label = res.groupdict()["label"] if labeled else -1
-      yield {
+      record = {
           "text": text,
           "label": label,
       }
+      if keys and self.version.implements(tfds.core.Experiment.S3):
+        yield path, record
+      else:
+        yield record

tensorflow_datasets/text/imdb_test.py

@@ -33,5 +33,9 @@ class IMDBReviewsTest(testing.DatasetBuilderTestCase):
   DL_EXTRACT_RESULT = "aclImdb_v1.tar.gz"
 
 
+class IMDBReviewsS3Test(IMDBReviewsTest):
+  VERSION = "experimental_latest"
+
+
 if __name__ == "__main__":
   testing.test_main()