Repository: huggingface/datasets
Pull request: Support streaming for pubmed
This PR makes some minor changes to the `pubmed` dataset to allow for `streaming=True`. Fixes #3739.

Basically, I followed the C4 dataset which works in streaming mode as an example, and made the following changes:

* Change URL prefix from `ftp://` to `https://`
* Explicilty `open` the filename and pass the XML contents to `etree.fromstring(xml_str)`

The Github diff tool makes it look like the changes are larger than they are, sorry about that.

I tested locally and the `pubmed` dataset now works in both normal and streaming modes. There is some overhead at the start of each shard in streaming mode as building the XML tree online is quite slow (each pubmed .xml.gz file is ~20MB), but the overhead gets amortized over all the samples in the shard. On my laptop with a single CPU worker I am able to stream at about ~600 samples/s.

Problem statement:
Pubmed dataset does not work in streaming mode
## Describe the bug
Trying to use the `pubmed` dataset with `streaming=True` fails.

## Steps to reproduce the bug
```python
import datasets
pubmed_train = datasets.load_dataset('pubmed', split='train', streaming=True)
print (next(iter(pubmed_train)))
```

## Expected results
I would expect to see the first training sample from the pubmed dataset. 

## Actual results
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/abhinav/Documents/mosaicml/mosaicml_venv/lib/python3.8/site-packages/datasets/iterable_dataset.py", line 367, in __iter__
    for key, example in self._iter():
  File "/Users/abhinav/Documents/mosaicml/mosaicml_venv/lib/python3.8/site-packages/datasets/iterable_dataset.py", line 364, in _iter
    yield from ex_iterable
  File "/Users/abhinav/Documents/mosaicml/mosaicml_venv/lib/python3.8/site-packages/datasets/iterable_dataset.py", line 79, in __iter__
    for key, example in self.generate_examples_fn(**self.kwargs):
  File "/Users/abhinav/.cache/huggingface/modules/datasets_modules/datasets/pubmed/9715addf10c42a7877a2149ae0c5f2fddabefc775cd1bd9b03ac3f012b86ce46/pubmed.py", line 373, in _generate_examples
    tree = etree.parse(filename)
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/xml/etree/ElementTree.py", line 1202, in parse
    tree.parse(source, parser)
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/xml/etree/ElementTree.py", line 584, in parse
    source = open(source, "rb")
FileNotFoundError: [Errno 2] No such file or directory: 'gzip://pubmed21n0001.xml::ftp://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed21n0001.xml.gz'
```
## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 1.18.2
- Platform: macOS-11.4-x86_64-i386-64bit
- Python version: 3.8.2
- PyArrow version: 6.0.0

## Comments
The error looks like an issue with `open` vs. `xopen` inside the `xml` package. It looks like it's trying to open the remote source URL, which has been edited with prefix `gzip://...`. 

Maybe there can be an explicit `xopen` before passing the raw data to `etree`, something like:

```python
# Before
tree = etree.parse(filename)
root = tree.getroot()

# After
with xopen(filename) as f:
  data_str = f.read()
root = etree.fromstring(data_str)
```

Hints from issue discussion:

Commit message: update hardcoded exception

Patch to review:
diff --git a/datasets/pubmed/pubmed.py b/datasets/pubmed/pubmed.py
index 4f12617f6ff..4d71bb58e7d 100644
--- a/datasets/pubmed/pubmed.py
+++ b/datasets/pubmed/pubmed.py
@@ -39,7 +39,7 @@
 # The HuggingFace dataset library don't host the datasets but only point to the original files
 # This can be an arbitrary nested dict/list of URLs (see below in `_split_generators` method)
 # Note these URLs here are used by MockDownloadManager.create_dummy_data_list
-_URLs = [f"ftp://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed22n{i:04d}.xml.gz" for i in range(1, 1115)]
+_URLs = [f"https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed22n{i:04d}.xml.gz" for i in range(1, 1115)]
 
 
 # Copyright Ferry Boender, released under the MIT license.
@@ -358,28 +358,29 @@ def _generate_examples(self, filenames):
         """Yields examples."""
         id_ = 0
         for filename in filenames:
-            try:
-                tree = etree.parse(filename)
-                root = tree.getroot()
-                xmldict = self.xml_to_dictionnary(root)
-            except etree.ParseError:
-                logger.warning(f"Ignoring file {filename}, it is malformed")
-                continue
-
-            for article in xmldict["PubmedArticleSet"]["PubmedArticle"]:
-                self.update_citation(article)
-                new_article = default_article()
-
+            with open(filename, "rt", encoding="utf-8") as f:
+                xml_str = f.read()
                 try:
-                    deepupdate(new_article, article)
-                except Exception:
-                    logger.warning(f"Ignoring article {article}, it is malformed")
+                    root = etree.fromstring(xml_str)
+                    xmldict = self.xml_to_dictionnary(root)
+                except etree.ParseError:
+                    logger.warning(f"Ignoring file {filename}, it is malformed")
                     continue
 
-                try:
-                    _ = self.info.features.encode_example(new_article)
-                except Exception as e:
-                    logger.warning(f"Ignore example because {e}")
-                    continue
-                yield id_, new_article
-                id_ += 1
+                for article in xmldict["PubmedArticleSet"]["PubmedArticle"]:
+                    self.update_citation(article)
+                    new_article = default_article()
+
+                    try:
+                        deepupdate(new_article, article)
+                    except Exception:
+                        logger.warning(f"Ignoring article {article}, it is malformed")
+                        continue
+
+                    try:
+                        _ = self.info.features.encode_example(new_article)
+                    except Exception as e:
+                        logger.warning(f"Ignore example because {e}")
+                        continue
+                    yield id_, new_article
+                    id_ += 1
diff --git a/src/datasets/utils/mock_download_manager.py b/src/datasets/utils/mock_download_manager.py
index fcae4665f47..8caa5a1324f 100644
--- a/src/datasets/utils/mock_download_manager.py
+++ b/src/datasets/utils/mock_download_manager.py
@@ -172,7 +172,7 @@ def create_dummy_data_list(self, path_to_dummy_data, data_url):
         # trick: if there are many shards named like `data.txt-000001-of-00300`, only use the first one
         is_tf_records = all(bool(re.findall("[0-9]{3,}-of-[0-9]{3,}", url)) for url in data_url)
         is_pubmed_records = all(
-            url.startswith("ftp://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed") for url in data_url
+            url.startswith("https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed") for url in data_url
         )
         if is_tf_records or is_pubmed_records:
             data_url = [data_url[0]] * len(data_url)


Code review comments:
diff_hunk: @@ -358,28 +358,29 @@ def _generate_examples(self, filenames):
         """Yields examples."""
         id_ = 0
         for filename in filenames:
-            try:
-                tree = etree.parse(filename)
-                root = tree.getroot()
-                xmldict = self.xml_to_dictionnary(root)
-            except etree.ParseError:
-                logger.warning(f"Ignoring file {filename}, it is malformed")
-                continue
-
-            for article in xmldict["PubmedArticleSet"]["PubmedArticle"]:
-                self.update_citation(article)
-                new_article = default_article()
-
+            with open(filename, "rt", encoding="utf-8") as f:
+                xml_str = f.read()
                 try:
-                    deepupdate(new_article, article)
-                except Exception:
-                    logger.warning(f"Ignoring article {article}, it is malformed")
+                    root = etree.fromstring(xml_str)
+                    xmldict = self.xml_to_dictionnary(root)
+                except etree.ParseError:
+                    logger.warning(f"Ignoring file {filename}, it is malformed")
                     continue
 
-                try:
-                    _ = self.info.features.encode_example(new_article)
-                except Exception as e:
-                    logger.warning(f"Ignore example because {e}")
-                    continue
-                yield id_, new_article
-                id_ += 1
+                for article in xmldict["PubmedArticleSet"]["PubmedArticle"]:
+                    self.update_citation(article)
+                    new_article = default_article()
+
+                    try:
+                        deepupdate(new_article, article)
+                    except Exception:
+                        logger.warning(f"Ignoring article {article}, it is malformed")
+                        continue
+
+                    try:
+                        _ = self.info.features.encode_example(new_article)
+                    except Exception as e:
+                        logger.warning(f"Ignore example because {e}")
+                        continue
+                    yield id_, new_article
+                    id_ += 1
path: datasets/pubmed/pubmed.py
text: @user1:
I don't think this is necessary, once that `xml.etree.ElementTree.parse` already supports streaming. See:
- #3476

The only thing needs being changed is the import: `as ET` instead of `as etree`: https://github.com/huggingface/datasets/blob/71aa41caee3d2e9fde981cacc4a5fbddec58250d/datasets/pubmed/pubmed.py#L19