Support streaming for pubmed
Opened by abhi-mosaic (closed)

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.

Comments:
@albertvillanova just FYI, since you were so helpful with the previous pubmed issue :) 

IIRC streaming from FTP is not fully tested yet, so I'm fine with switching to HTTPS for now, as long as the download speed/availability is great

@albertvillanova Thanks for pointing me to the `ET` module replacement. It should look a lot cleaner now.

Unfortunately I tried keeping the `ftp://` protocol but was seeing timeout errors? in streaming mode (below). I think the `https://` performance is not an issue, when I was profiling the `open(..) -> f.read() -> etree.fromstring(xml_str)` codepath, most of the time was spent in the XML parsing rather than the data download.


Error when using `ftp://`:

```
Traceback (most recent call last):
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/venv/lib/python3.8/site-packages/fsspec/implementations/ftp.py", line 301, in _fetch_range
    self.fs.ftp.retrbinary(
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/ftplib.py", line 430, in retrbinary
    callback(data)
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/venv/lib/python3.8/site-packages/fsspec/implementations/ftp.py", line 293, in callback
    raise TransferDone
fsspec.implementations.ftp.TransferDone

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "test_pubmed_streaming.py", line 9, in <module>
    print (next(iter(pubmed_train_streaming)))
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/abhi-datasets/src/datasets/iterable_dataset.py", line 365, in __iter__
    for key, example in self._iter():
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/abhi-datasets/src/datasets/iterable_dataset.py", line 362, in _iter
    yield from ex_iterable
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/abhi-datasets/src/datasets/iterable_dataset.py", line 79, in __iter__
    yield from self.generate_examples_fn(**self.kwargs)
  File "/Users/abhinav/.cache/huggingface/modules/datasets_modules/datasets/pubmed/af552ed918e2841e8427203530e3cfed3a8bc3213041d7853bea1ca67eec683d/pubmed.py", line 362, in _generate_examples
    tree = ET.parse(filename)
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/abhi-datasets/src/datasets/streaming.py", line 65, in wrapper
    return function(*args, use_auth_token=use_auth_token, **kwargs)
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/abhi-datasets/src/datasets/utils/streaming_download_manager.py", line 636, in xet_parse
    return ET.parse(f, parser=parser)
  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 595, in parse
    self._root = parser._parse_whole(source)
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/abhi-datasets/src/datasets/utils/streaming_download_manager.py", line 293, in read_with_retries
    out = read(*args, **kwargs)
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/gzip.py", line 292, in read
    return self._buffer.read(size)
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/_compression.py", line 68, in readinto
    data = self.read(len(byte_view))
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/gzip.py", line 479, in read
    if not self._read_gzip_header():
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/gzip.py", line 422, in _read_gzip_header
    magic = self._fp.read(2)
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/gzip.py", line 96, in read
    self.file.read(size-self._length+read)
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/venv/lib/python3.8/site-packages/fsspec/spec.py", line 1485, in read
    out = self.cache._fetch(self.loc, self.loc + length)
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/venv/lib/python3.8/site-packages/fsspec/caching.py", line 153, in _fetch
    self.cache = self.fetcher(start, end)  # new block replaces old
  File "/Users/abhinav/Documents/mosaicml/hf_datasets/venv/lib/python3.8/site-packages/fsspec/implementations/ftp.py", line 311, in _fetch_range
    self.fs.ftp.getmultiline()
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/ftplib.py", line 224, in getmultiline
    line = self.getline()
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/ftplib.py", line 206, in getline
    line = self.file.readline(self.maxline + 1)
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/socket.py", line 669, in readinto
    return self._sock.recv_into(b)
socket.timeout: timed out
```