unify shuff typo

wenet-e2e · Mar 20, 2024 · 0f2f339 · 0f2f339
1 parent 11eb863
commit 0f2f339
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/wenet/dataset/dataset.py b/wenet/dataset/dataset.py
@@ -42,9 +42,13 @@ def Dataset(data_type,
     assert data_type in ['raw', 'shard']
     # cycle dataset
     cycle = conf.get('cycle', 1)
-    list_shuffle = conf.get('list_shuffle', False)
-    list_shuffle_size = conf.get('list_shuffle_size', 10000)
-
+    # stage1 shuffle: source
+    list_shuffle = conf.get('list_shuffle', True)
+    list_shuffle_size = 10000000
+    if list_shuffle:
+        list_shuffle_conf = conf.get('list_shuffle_conf', {})
+        list_shuffle_size = list_shuffle_conf.get('shuffle_size',
+                                                  list_shuffle_size)
     if data_type == 'raw':
         dataset = WenetRawDatasetSource(data_list_file,
                                         partition=partition,