Enabling SST2 dataset usage in fbcode

torchtext/_download_hooks.py

@@ -3,6 +3,10 @@
 from tqdm import tqdm
 # This is to allow monkey-patching in fbcode
 from torch.hub import load_state_dict_from_url # noqa
+from torchtext._internal.module_utils import is_module_available
+
+if is_module_available("torchdata"):
+    from torchdata.datapipes.iter import HttpReader # noqa F401
 
 
 def _stream_response(r, chunk_size=16 * 1024):

torchtext/experimental/datasets/sst2.py

@@ -10,10 +10,11 @@
 )
 
 if is_module_available("torchdata"):
-    from torchdata.datapipes.iter import (
-        HttpReader,
-        IterableWrapper,
-    )
+    from torchdata.datapipes.iter import IterableWrapper
+    # we import HttpReader from _download_hooks so we can swap out public URLs
+    # with interal URLs when the dataset is used within Facebook
+    from torchtext._download_hooks import HttpReader
+
 
 NUM_LINES = {
     "train": 67349,