fix caching

Author: parmeet

torchtext/datasets/amazonreviewpolarity.py

@@ -23,8 +23,8 @@
 _PATH = 'amazon_review_polarity_csv.tar.gz'
 
 _EXTRACTED_FILES = {
-    'train': os.path.join(_PATH, 'amazon_review_polarity_csv', 'train.csv'),
-    'test': os.path.join(_PATH, 'amazon_review_polarity_csv', 'test.csv'),
+    'train': os.path.join('amazon_review_polarity_csv', 'train.csv'),
+    'test': os.path.join('amazon_review_polarity_csv', 'test.csv'),
 }
 
 
@@ -44,9 +44,25 @@ def AmazonReviewPolarity(root: str, split: Union[Tuple[str], str]):
         filepath_fn=lambda x: os.path.join(root, _PATH), hash_dict={os.path.join(root, _PATH): MD5}, hash_type="md5"
     )
     cache_compressed_dp = GDriveReader(cache_compressed_dp).end_caching(mode="wb", same_filepath_fn=True)
+
+    def extracted_filepath_fn(x):
+        file_path = os.path.join(root, _EXTRACTED_FILES[split])
+        dir_path = os.path.dirname(file_path)
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+        return file_path
+
     cache_decompressed_dp = cache_compressed_dp.on_disk_cache(
-        filepath_fn=lambda x: os.path.join(root, os.path.dirname(_EXTRACTED_FILES[split]), os.path.basename(x)))
-    cache_decompressed_dp = FileOpener(cache_decompressed_dp, mode="b").read_from_tar()
-    cache_compressed_dp = cache_decompressed_dp.end_caching(mode="wb", same_filepath_fn=True)
-    data_dp = FileOpener(cache_decompressed_dp.filter(lambda x: _EXTRACTED_FILES[split] in x[0]).map(lambda x: x[0]), mode='b')
+        filepath_fn=extracted_filepath_fn)
+    cache_decompressed_dp = FileOpener(cache_decompressed_dp, mode="b").\
+        read_from_tar().\
+        filter(lambda x: _EXTRACTED_FILES[split] in x[0]).\
+        map(lambda x: (x[0].replace('_PATH' + '/', ''), x[1]))
+    cache_decompressed_dp = cache_decompressed_dp.end_caching(mode="wb", same_filepath_fn=True)
+    data_dp = FileOpener(cache_decompressed_dp, mode='b')
+
+    # data_dp = FileOpener(cache_compressed_dp, mode='b')
+    # data_dp = data_dp.read_from_tar()
+    # data_dp = data_dp.filter(lambda x: _EXTRACTED_FILES[split] in x[0])
+
     return data_dp.parse_csv().map(fn=lambda t: (int(t[0]), ' '.join(t[1:])))