fix error when read shard

main.py

@@ -73,46 +73,49 @@ def prepare_dataset(batch, processor):
 
 
 def load_prepared_dataset(path, processor, cache_file_filter_name, cache_file_map_name, num_proc=5):
-    dataset = load_from_disk(path)
-    list_cache_prefetch_files = glob.glob(
-        cache_file_map_name.replace(cache_processing_dataset_folder, cache_processing_dataset_folder_prefetch).replace(
-            '.arrow', '*'))
-
-    # Do not re-compute what already in cache folder
-    if cache_file_map_name.startswith(cache_processing_dataset_folder_prefetch):
-        if len(glob.glob(cache_file_map_name.replace(cache_processing_dataset_folder_prefetch,
-                                                 cache_processing_dataset_folder).replace('.arrow', '*'))) > 0:
-            return
-        if len(list_cache_prefetch_files) > 0:
-            return
-
-    # check cache file
-    if len(glob.glob(cache_file_map_name.replace('.arrow', '*'))) == 0 and len(list_cache_prefetch_files) > 0:
-        for item_file in list_cache_prefetch_files:
-            shutil.move(item_file, item_file.replace(cache_processing_dataset_folder_prefetch,
-                                                   cache_processing_dataset_folder))
-    if len(glob.glob(cache_file_map_name.replace('.arrow', '*'))) > 0:
-        return dataset.map(prepare_dataset,
-                           remove_columns=dataset.column_names,
-                           batch_size=32,
-                           num_proc=num_proc,
-                           batched=True,
-                           fn_kwargs={"processor": processor},
-                           cache_file_name=cache_file_map_name)
-
-    dataset = dataset.filter(lambda example: len(example['speech']) < 160000,
-                             batch_size=32,
-                             num_proc=num_proc,
-                             cache_file_name=cache_file_filter_name)
-    processed_dataset = dataset.map(prepare_dataset,
-                                    remove_columns=dataset.column_names,
-                                    batch_size=32,
-                                    num_proc=num_proc,
-                                    batched=True,
-                                    fn_kwargs={"processor": processor},
-                                    cache_file_name=cache_file_map_name)
-    processed_dataset.cleanup_cache_files()
-    return processed_dataset
+    try:
+      dataset = load_from_disk(path)
+      list_cache_prefetch_files = glob.glob(
+          cache_file_map_name.replace(cache_processing_dataset_folder, cache_processing_dataset_folder_prefetch).replace(
+              '.arrow', '*'))
+
+      # Do not re-compute what already in cache folder
+      if cache_file_map_name.startswith(cache_processing_dataset_folder_prefetch):
+          if len(glob.glob(cache_file_map_name.replace(cache_processing_dataset_folder_prefetch,
+                                                  cache_processing_dataset_folder).replace('.arrow', '*'))) > 0:
+              return
+          if len(list_cache_prefetch_files) > 0:
+              return
+
+      # check cache file
+      if len(glob.glob(cache_file_map_name.replace('.arrow', '*'))) == 0 and len(list_cache_prefetch_files) > 0:
+          for item_file in list_cache_prefetch_files:
+              shutil.move(item_file, item_file.replace(cache_processing_dataset_folder_prefetch,
+                                                    cache_processing_dataset_folder))
+      if len(glob.glob(cache_file_map_name.replace('.arrow', '*'))) > 0:
+          return dataset.map(prepare_dataset,
+                            remove_columns=dataset.column_names,
+                            batch_size=32,
+                            num_proc=num_proc,
+                            batched=True,
+                            fn_kwargs={"processor": processor},
+                            cache_file_name=cache_file_map_name)
+
+      dataset = dataset.filter(lambda example: len(example['speech']) < 160000,
+                              batch_size=32,
+                              num_proc=num_proc,
+                              cache_file_name=cache_file_filter_name)
+      processed_dataset = dataset.map(prepare_dataset,
+                                      remove_columns=dataset.column_names,
+                                      batch_size=32,
+                                      num_proc=num_proc,
+                                      batched=True,
+                                      fn_kwargs={"processor": processor},
+                                      cache_file_name=cache_file_map_name)
+      processed_dataset.cleanup_cache_files()
+      return processed_dataset
+    except:
+      return None
 
 
 def commit_checkpoint():
@@ -264,6 +267,9 @@ if __name__ == "__main__":
                                                                               'cache-test-map-shard-{}.arrow'.format(
                                                                                   test_dataset_shard_idx))
                                              )
+        if train_dataset is None or test_dataset is None:
+          print("Ignore Shard {}".format(train_dataset_shard_idx))
+          continue
         test_dataset = test_dataset.shard(num_test_sub_shard, idx_sub_shard)
 
         # Prefetch_dataset