Repository: pytorch/text

Issue:
Migrate datasets to build on top of torchdata datapipes
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

**Motivation**

https://github.com/pytorch/data#why-composable-data-loading

**_user-experience:_** TorchData datasets enable new functional API, auto-sharding, and snapshotting support out-of-the-box. They also enable standard flow-control like batching, collation, shuffling, bucketing, and mapping/transformation using user-defined functions and transforms (UDFs).

**_Maintenance:_** By relying on TorchData, we no longer have to maintain low level functionality like downloading, extracting, caching, file/steam parsing, etc.

**Reference** 
Examples: https://github.com/facebookexternal/torchdata/tree/main/examples/text
TorchData: https://github.com/facebookexternal/torchdata

## Backlog of datasets

- [x] AG_NEWS https://github.com/pytorch/text/pull/1498
- [x] AmazonReviewFull https://github.com/pytorch/text/pull/1499
- [x] AmazonReviewPolarity https://github.com/pytorch/text/pull/1490
- [x] DBpedia https://github.com/pytorch/text/pull/1500
- [x] SogouNews https://github.com/pytorch/text/pull/1503
- [x] YelpReviewFull https://github.com/pytorch/text/pull/1507
- [x] YelpReviewPolarity https://github.com/pytorch/text/pull/1509
- [x] YahooAnswers https://github.com/pytorch/text/pull/1508
- [x] CoNLL2000Chunking https://github.com/pytorch/text/pull/1515
- [x] UDPOS https://github.com/pytorch/text/pull/1535
- [x] IWSLT2016 https://github.com/pytorch/text/pull/1545
- [x] IWSLT2017 https://github.com/pytorch/text/pull/1547
- [x] Multi30K https://github.com/pytorch/text/pull/1536
- [x] SQuAD1 https://github.com/pytorch/text/pull/1513
- [x] SQuAD2 https://github.com/pytorch/text/pull/1514
- [x] PennTreebank https://github.com/pytorch/text/pull/1511
- [x] WikiText103 https://github.com/pytorch/text/pull/1518
- [x] WikiText2 https://github.com/pytorch/text/pull/1519
- [x] EnWik9 https://github.com/pytorch/text/pull/1512
- [x] IMDB https://github.com/pytorch/text/pull/1531
- [x] SST2 https://github.com/pytorch/text/pull/1538
- [x] CC-100 https://github.com/pytorch/text/pull/1562


## Contributing
Please leave a message below if you plan to work on particular dataset(s) to avoid duplication of efforts. Also please link to the corresponding PRs.

cc: @Nayef211 , @abhinavarora , @erip , @ejguan , @VitalyFedyunin 



Discussion:
This sounds great! Logistically, how do we want to organize the PRs? Should we have a branch which adds torchdata as a dependency, we branch from that branch and target that branch for our PRs?
@erip That's a good point. Right now we have torchdata as an optional dependency since it is not yet available on pypi (torchdata is in prototype phase until the next Pytorch release). I think going forward we still want to keep it as an optional dependency until torchdata is released as beta.

What do you think @parmeet @abhinavarora?
Right, agreed with @Nayef211 . Let's continue keeping optional dependency. I am working on PR to migrate AmazonReviewPolarity (https://github.com/pytorch/text/pull/1490) to ensure nothing breaks in the CI when doing the migration. Once this lands, we can follow the same to migrate other datasets. 
Great, so individual PRs instead of one big one -- got it! I'll pick a couple soon and start digging.
I guess it will become necessary to use some `skip if` logic when dealing with optional dependency.
In CI, I recommend to make them fail unless explicitly allowed to skip.
We had an issue where CUDA tests were unintentionally skipped because of wrong CI configuration and permissive `skipIf`.

Here is the update I applied recently to avoid that https://github.com/pytorch/audio/pull/2127
I'll keep a moving list here. To start, I'll take...

- AG_NEWS
- Amazon review full
- DBPedia
- Sogou News
- YelpReviewFull
- YelpReviewPolarity
- YahooAnswers 
- CoNLL2000Chunking
- SQUAD1
- SQUAD2
- UDPOS (blocked by CoNLL2000Chunking)
- IWSLT2016
- IWSLT2017
- Multi30K
Hey @erip, thank you for picking the datasets :). 

> I'll keep a moving list here. To start, I'll take...

Since we already have a backlog checklist in the the issue, you may constraint this list to the datasets you pick-up and checkmark it directly in the backlog items once the PR lands.

> * ~AmazonReviewPolarity~ Parmeet has finished this

 you may skip the mention for this dataset. I will check mark it once I land the PR :)





AmazonRiviewPolarity  [WIP] https://github.com/pytorch/text/pull/1490. Note that this PR shall also resolve issues with existing CI dataset tests. 
> I guess it will become necessary to use some `skip if` logic when dealing with optional dependency. In CI, I recommend to make them fail unless explicitly allowed to skip. We had an issue where CUDA tests were unintentionally skipped because of wrong CI configuration and permissive `skipIf`.
> 
> Here is the update I applied recently to avoid that [pytorch/audio#2127](https://github.com/pytorch/audio/pull/2127)

Actually in this case, the problem is beyond just unit-testing. The package builds (conda, wheels, docs) also rely on availability of torchdata, because it is imported inside the modules. As of now, the import is [conditioned](https://github.com/pytorch/text/blob/1a052693509ce32b6fb91302f6ca62546b0afe0d/torchtext/experimental/datasets/sst2.py#L12) on availability. For unit-testing, we are [installing torchdata](https://github.com/pytorch/text/blob/1a052693509ce32b6fb91302f6ca62546b0afe0d/.circleci/unittest/linux/scripts/install.sh#L17) from source. Perhaps we could do the same for package builds?
> Actually in this case, the problem is beyond just unit-testing. The package builds (conda, wheels, docs) also rely on availability of torchdata, because it is imported inside the modules. As of now, the import is [conditioned](https://github.com/pytorch/text/blob/1a052693509ce32b6fb91302f6ca62546b0afe0d/torchtext/experimental/datasets/sst2.py#L12) on availability. For unit-testing, we are [installing torchdata](https://github.com/pytorch/text/blob/1a052693509ce32b6fb91302f6ca62546b0afe0d/.circleci/unittest/linux/scripts/install.sh#L17) from source. Perhaps we could do the same for package builds?

Isn't torchdata optional runtime dependency? I looked into the build process, and CI configuration but I do not see any mention to torchdata or the build/packaging logic changed by the presence of torchdata.
I plan to pick-up CC-100 next..
> Isn't torchdata optional runtime dependency? 

Yes, it is if users does not use migrated datasets. But they would need installation if they need to use them since the underlying implementation of migrated datasets does not fall back to previous (non-datapipes based) implementation if torchdata package is not found. 

Edit: We are going to make torchdata hard-dependency with the next release. I think the question is rather what should we be doing during this migration phase. As a safe side, I think keeping it optional make sense, so that users working directly from source/nightly do not break their workflow unless they depend on the migrated datasets (in which case they would need installation anyway).


> I looked into the build process, and CI configuration but I do not see any mention to torchdata or the build/packaging logic changed by the presence of torchdata.

If we remove the [conditional import](https://github.com/pytorch/text/blob/1a052693509ce32b6fb91302f6ca62546b0afe0d/torchtext/experimental/datasets/sst2.py#L12), it breaks the CI build/packaging related tests. I guess we would need to install torchdata in appropriate locations for them to pass like we are doing for unit-testing [here](https://github.com/pytorch/text/blob/1a052693509ce32b6fb91302f6ca62546b0afe0d/.circleci/unittest/linux/scripts/install.sh#L17)


> Edit: We are going to make torchdata hard-dependency with the next release. I think the question is rather what should we be doing during this migration phase. As a safe side, I think keeping it optional make sense, so that users working directly from source/nightly do not break their workflow unless they depend on the migrated datasets (in which case they would need installation anyway).

I second on this. Before our release, TorchData theoretically is unstable library, which means our API would change a lot. Nightly users of TorchText can encounter error during installation phase. This should not be expected as some users won't use anything from TorchData.
> Edit: We are going to make torchdata hard-dependency with the next release.

Okay, this is what I wanted to verify. As long as it is the intended move (and is going to be mentioned in README) then it sounds good to me.

> If we remove the conditional import, it breaks the CI build/packaging related tests. I guess we would need to install torchdata in appropriate locations for them to pass like we are doing for unit-testing here

nit: IMU, packaging job should not be importing torchtext, so, I am not entirely convinced that it's required at build time. (especially since torchdata is pure-python package) Maybe it failed at smoke test step that happens right after the packaging step in the same CI job.

> nit: IMU, packaging job should not be importing torchtext, so, I am not entirely convinced that it's required at build time. (especially since torchdata is pure-python package) Maybe it failed at smoke test step that happens right after the packaging step in the same CI job.

sorry for the confusion, yes you are right it is not needed for packaging, but rather in testing in the same CI job (basically importing torchtext would fail as it would implicitly import datasets that would fail due to missing torchdata package)

I plan to pick up 

- WikiText103
- WikiText2
Continuing on the discussion (https://github.com/pytorch/text/pull/1513#issuecomment-1013319134) by @ejguan, I think it's worth considering this point. The point of discussion is whether we should keep both the implementations (current one and the one with datapipes) in upcoming release (0.12) and  provide user-warning to switch to the new implementation? 

**Some thoughts**

Pros:
* Better handling of BC related issues, risk minimization of breaking user workflows in unexpected ways
* Give enough time to users to switch to new implementation, meanwhile making some educational blogs, tutorial for transition 
* Minimize any risk related to torchdata being still in process of making first beta-release or API changes/adaptions

Cons:
* Cannot think of any, except a bit of extra maintenance


@mthrok, @Nayef211 , @abhinavarora, @erip:  would appreciate your thoughts here to help us reach a consensus. 


 > Pros:
> 
> * Better handling of BC related issues, risk minimization of breaking user workflows in unexpected ways
> * Give enough time to users to switch to new implementation, meanwhile making some educational blogs, tutorial for transition
> * Minimize any risk related to torchdata being still in process of making first beta-release or API changes/adaptions

> Cons:
> 
> * Cannot think of any, except a bit of extra maintenance

In terms of pros I agree with the 2nd and 3rd points. My question about the first point would be whether there would be any BC related issues? As we discussed before starting this revamp effort, the purpose of this migration is to change the underlying implementation of the datasets without modifying the user facing API. If this is indeed the case, the end-user shouldn't notice any changes in how they are using the dataset. The main issues I could foresee coming up would be related to dependency changes caused by upgrading to datapipes. 


That being said, I am still okay with maintaining both implementations if that would guarantee a better UX. Would we want to add all the datapipe backed datasets in `experimental/datasets` folder if we ended up maintaining both implementations?



> Continuing on the discussion ([#1513 (comment)](https://github.com/pytorch/text/pull/1513#issuecomment-1013319134)) by @ejguan, I think it's worth considering this point. The point of discussion is whether we should keep both the implementations (current one and the one with datapipes) in upcoming release (0.12) and provide user-warning to switch to the new implementation?
> 
> **Some thoughts**
> 
> Pros:
> 
> * Better handling of BC related issues, risk minimization of breaking user workflows in unexpected ways
> * Give enough time to users to switch to new implementation, meanwhile making some educational blogs, tutorial for transition
> * Minimize any risk related to torchdata being still in process of making first beta-release or API changes/adaptions
> 
> Cons:
> 
> * Cannot think of any, except a bit of extra maintenance
> 
> @mthrok, @Nayef211 , @abhinavarora, @erip: would appreciate your thoughts here to help us reach a consensus.

I would be wary of adding additional maintenance costs. Could you elaborate what is the BC breaking change. Since the API is the same, is it just the missing torchdata library? Given that torchdata is also releasing with our 0.12 release, I honestly don't see any serious issue with just migrating to the new code. Also the test failures will only impact nightly users right? In that case the warning message is very straightforward and the users can easily set up torchdata. Lease let me know if I am missing something. 
Thanks @abhinavarora  and @Nayef211  for your thoughts on BC issue. Honestly, I am of similar impression and cannot think immediately of any obvious BC issues (which is why I wrote unexpected :) ). One of things new API doesn't support is `__next__` method, though I don't know how this might impact user if any.

@Nayef211  regarding you question: 
> Would we want to add all the datapipe backed datasets in experimental/datasets folder if we ended up maintaining both implementations?

I don't think we want to do that. One way to keep both the implementations is by switching to new implementation through some flag etc. By default users would still use the existing implementation and along with that get user-warning and settings information to switch to new implementation. This surely adds to some maintenance overhead which @abhinavarora already raised concerns about.




I'd like to take a stab at the IMDB dataset
Replying the comment from #1513

> > BTW, do we want to replace the old implementation directly? I think we may want to keep the original implementation at least one release cycle.
> 
> Thanks @ejguan for this call out. I think you made a valid point. Perhaps we could support both the implementation with the ability to switch to new implementation by setting flags or something. Let's take this discussion on the issue (#1494 ) directly.

If using a flag, I recommend that to make it a part of constructor. (should not be a big deal if the interface is same) Not a global flag. torchaudio has a global flag for I/O backend and it's not a good maintenance experience. It turned the codebase stateful and in a large code base, it becomes impossible to tell which downstream project sets/expects which flag state.

Quick update:

- UDPOS is done. I can't check if off here, but maybe someone else can. 😄 
- The IWSLT datasets are a bit tricky for a couple of reasons:
  - The have nested tarballs. I have submitted a PR upstream to torchdata to add a [flatmap datapipe](https://github.com/pytorch/data/pull/179) and plan to add it to the torchtext repo until that's merged. 
  - The files get cleaned and rewritten in-place. This isn't a strict requirement, but it's a bit delicate, so I'm trying to proceed with caution.
  - There's a lot of conditional logic depending on the file format. I'll need to think about how best to handle this. Probably just mapping a function which abstracts the conditions... hope it'll be serializable

In the meantime, I'll take a stab at Multi30k -- it looks quite straightforward. I can't seem to find CC100 and SST2 impls -- are these new datasets for torchtext?
> * UDPOS is done. I can't check if off here, but maybe someone else can. 😄

Done :) 
  
>   * The have nested tarballs. I have submitted a PR upstream to torchdata to add a [flatmap datapipe](https://github.com/pytorch/data/pull/179) and plan to add it to the torchtext repo until that's merged.

Could you please help provide more details as to where is this required?


> In the meantime, I'll take a stab at Multi30k

Great, thanks!

> I can't seem to find CC100 and SST2 impls -- are these new datasets for torchtext?

SST-2 is in experimental folder. @Nayef211  is working on it, and is currently blocked due to some internal fbcode dependency. I am looking at CC100. It is already added as a notebook example [here](https://github.com/pytorch/data/blob/main/examples/text/CC100.ipynb). Unfortunately, the host seems to be quite un-responsive and hence i am currently blocked to add this one. I am trying to contact some internal folks to figure out the path forward.



> Could you please help provide more details ...

Yes, it's [here](https://github.com/pytorch/text/blob/main/torchtext/datasets/iwslt2016.py#L231-L241). We extract contents from a tarball and inside the tarball are other tarballs. This requires something roughly like:

```python
# [ 1.tgz, 2.tgz, ..., n.tgz]
inner_tars = outer_tar.read_from_tar()

# Yuck, an IterableDataPipe of IterableDataPipes... flatmap to the rescue
# [ 1_1.txt, 1_2.txt, 2_1.txt, 2_2.txt, 3_1.txt, ...]
inner_inner_files = inner_tars.flatmap(lambda inner_tar: FileOpener(inner_tar).read_from_tar())
```
> inner_inner_files = inner_tars.flatmap(lambda inner_tar: FileOpener(inner_tar).read_from_tar())

I see, but isn't it that tar datapipe would recursively already do it and return all the inner files? 

@parmeet oh, I'm not sure... I didn't think so, but I guess I hadn't tried yet so I don't know. I'll give it a shot tomorrow. 
@parmeet looking at the implementation of the datapipe associated with the `read_from_tar` functional, it doesn't extract recursively. The PR in torchdata seems to moving forward so it will likely not be something we need to worry about much.
> The PR in torchdata seems to moving forward so it will likely not be something we need to worry about much.

Awesome! Thanks so much @erip for taking it forward. I guess these are the only two datasets left (apart for CC-100 :) ). Please feel free to post if there are any blockages :). 



Resolving patch:
diff --git a/torchtext/datasets/__init__.py b/torchtext/datasets/__init__.py
--- a/torchtext/datasets/__init__.py
+++ b/torchtext/datasets/__init__.py
@@ -3,6 +3,7 @@
 from .ag_news import AG_NEWS
 from .amazonreviewfull import AmazonReviewFull
 from .amazonreviewpolarity import AmazonReviewPolarity
+from .cc100 import CC100
 from .conll2000chunking import CoNLL2000Chunking
 from .dbpedia import DBpedia
 from .enwik9 import EnWik9
@@ -26,6 +27,7 @@
     "AG_NEWS": AG_NEWS,
     "AmazonReviewFull": AmazonReviewFull,
     "AmazonReviewPolarity": AmazonReviewPolarity,
+    "CC100": CC100,
     "CoNLL2000Chunking": CoNLL2000Chunking,
     "DBpedia": DBpedia,
     "EnWik9": EnWik9,
diff --git a/torchtext/datasets/cc100.py b/torchtext/datasets/cc100.py
new file mode 100644
--- /dev/null
+++ b/torchtext/datasets/cc100.py
@@ -0,0 +1,56 @@
+import os.path
+
+from typing import Union, Tuple
+
+from torchtext._internal.module_utils import is_module_available
+
+if is_module_available("torchdata"):
+    from torchdata.datapipes.iter import FileOpener, HttpReader, IterableWrapper
+
+from torchtext.data.datasets_utils import (
+    _create_dataset_directory,
+    _wrap_split_argument
+)
+
+URL = "http://data.statmt.org/cc-100/%s.txt.xz"
+
+VALID_CODES = {
+    "am", "ar", "as", "az", "be", "bg", "bn", "bn_rom", "br", "bs", "ca", "cs", "cy", "da", "de",
+    "el", "en", "eo", "es", "et", "eu", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gn", "gu",
+    "ha", "he", "hi", "hi_rom", "hr", "ht", "hu", "hy", "id", "ig", "is", "it", "ja", "jv", "ka",
+    "kk", "km", "kn", "ko", "ku", "ky", "la", "lg", "li", "ln", "lo", "lt", "lv", "mg", "mk", "ml",
+    "mn", "mr", "ms", "my", "my_zaw", "ne", "nl", "no", "ns", "om", "or", "pa", "pl", "ps", "pt",
+    "qu", "rm", "ro", "ru", "sa", "si", "sc", "sd", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv",
+    "sw", "ta", "ta_rom", "te", "te_rom", "th", "tl", "tn", "tr", "ug", "uk", "ur", "ur_rom", "uz",
+    "vi", "wo", "xh", "yi", "yo", "zh-Hans", "zh-Hant", "zu",
+}
+
+NUM_LINES = None
+MD5 = None
+
+DATASET_NAME = "CC100"
+
+
+@_create_dataset_directory(dataset_name=DATASET_NAME)
+@_wrap_split_argument(("train",))
+def CC100(root: str, split: Union[Tuple[str], str], language_code: str = "en"):
+    if language_code not in VALID_CODES:
+        raise ValueError(f"Invalid language code {language_code}")
+
+    url = URL % language_code
+    url_dp = IterableWrapper([url])
+    cache_compressed_dp = url_dp.on_disk_cache(
+        filepath_fn=lambda x: os.path.join(root, os.path.basename(url))
+    )
+
+    cache_compressed_dp = HttpReader(cache_compressed_dp)
+    cache_compressed_dp = cache_compressed_dp.end_caching(mode="wb", same_filepath_fn=True)
+
+    cache_decompressed_dp = cache_compressed_dp.on_disk_cache(
+        filepath_fn=lambda x: os.path.join(root, os.path.basename(x).rstrip(".xz"))
+    )
+    cache_decompressed_dp = FileOpener(cache_decompressed_dp, mode="b").read_from_xz()
+    cache_decompressed_dp = cache_decompressed_dp.end_caching(mode="wb")
+
+    data_dp = FileOpener(cache_decompressed_dp, mode="r").readlines(return_path=False)
+    return data_dp.map(lambda x: (language_code, x))


Test patch: