Fix the multi_news dataset.

PiperOrigin-RevId: 794084822

Author: The TensorFlow Datasets Authors

tensorflow_datasets/summarization/multi_news.py

@@ -16,8 +16,7 @@
 """Multi-News dataset."""
 
 import os
-
-from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
+from etils import epath
 import tensorflow_datasets.public_api as tfds
 
 _CITATION = """
@@ -42,7 +41,8 @@
   - summary: news summary.
 """
 
-_URL = "https://drive.google.com/uc?export=download&id=1vRY2wM6rlOZrf9exGTm5pXj5ExlVwJ0C"
+_URL_PATH = "https://huggingface.co/datasets/multi_news/resolve/main/data"
+
 
 _DOCUMENT = "document"
 _SUMMARY = "summary"
@@ -51,7 +51,7 @@
 class MultiNews(tfds.core.GeneratorBasedBuilder):
   """Multi-News dataset."""
 
-  VERSION = tfds.core.Version("1.0.0")
+  VERSION = tfds.core.Version("2.0.0")
 
   def _info(self):
     return tfds.core.DatasetInfo(
@@ -67,35 +67,35 @@ def _info(self):
 
   def _split_generators(self, dl_manager):
     """Returns SplitGenerators."""
-    extract_path = os.path.join(
-        dl_manager.download_and_extract(_URL), "multi-news-original"
-    )
-    return [
-        tfds.core.SplitGenerator(
-            name=tfds.Split.TRAIN,
-            gen_kwargs={"path": os.path.join(extract_path, "train")},
-        ),
-        tfds.core.SplitGenerator(
-            name=tfds.Split.VALIDATION,
-            gen_kwargs={"path": os.path.join(extract_path, "val")},
+    data_dict = {
+        "train_src": _URL_PATH + "train.src.cleaned",
+        "train_tgt": _URL_PATH + "train.tgt",
+        "val_src": _URL_PATH + "val.src.cleaned",
+        "val_tgt": _URL_PATH + "val.tgt",
+        "test_src": _URL_PATH + "test.src.cleaned",
+        "test_tgt": _URL_PATH + "test.tgt",
+    }
+    files = dl_manager.download_and_extract(data_dict)
+    return {
+        "train": self._generate_examples(
+            files["train_src"], files["train_tgt"]
         ),
-        tfds.core.SplitGenerator(
-            name=tfds.Split.TEST,
-            gen_kwargs={"path": os.path.join(extract_path, "test")},
+        "validation": self._generate_examples(
+            files["val_src"], files["val_tgt"]
         ),
-    ]
+        "test": self._generate_examples(files["test_src"], files["test_tgt"]),
+    }
 
-  def _generate_examples(self, path=None):
+  def _generate_examples(self, src_file, tgt_file):
     """Yields examples."""
-    with tf.io.gfile.GFile(
-        os.path.join(path + ".src")
-    ) as src_f, tf.io.gfile.GFile(os.path.join(path + ".tgt")) as tgt_f:
+    with epath.Path(src_file).open() as src_f, epath.Path(
+        tgt_file
+    ).open() as tgt_f:
       for i, (src_line, tgt_line) in enumerate(zip(src_f, tgt_f)):
         yield i, {
             # In original file, each line has one example and natural newline
             # tokens "\n" are being replaced with "NEWLINE_CHAR". Here restore
             # the natural newline token to avoid special vocab "NEWLINE_CHAR".
             _DOCUMENT: src_line.strip().replace("NEWLINE_CHAR", "\n"),
-            # Remove the starting token "- " for every target sequence.
-            _SUMMARY: tgt_line.strip().lstrip("- "),
+            _SUMMARY: tgt_line.strip().lstrip(),
         }

tensorflow_datasets/summarization/multi_news_test.py

@@ -26,8 +26,14 @@ class MultiNewsTest(testing.DatasetBuilderTestCase):
       "validation": 1,  # Number of fake validation example
       "test": 1,  # Number of fake test example
   }
-  DL_EXTRACT_RESULT = ""
-
+  DL_EXTRACT_RESULT = {
+      "train_src": "train.src.cleaned",
+      "train_tgt": "train.tgt",
+      "val_src": "val.src.cleaned",
+      "val_tgt": "val.tgt",
+      "test_src": "test.src.cleaned",
+      "test_tgt": "test.tgt",
+  }
 
 if __name__ == "__main__":
   testing.test_main()