fix shuffle in files

wenet-e2e · Sep 27, 2024 · 7723cb3 · 7723cb3
1 parent bf1664c
commit 7723cb3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/wenet/dataset/datapipes.py b/wenet/dataset/datapipes.py
@@ -452,6 +452,8 @@ def __init__(self,
  dp = IterableWrapperIterDataPipe(filenames)
  # 0 shard many jsonl files
  dp = dp.shuffle().repeat(cycle).shard(partition)
+ if shuffle:
+ self.dp = self.dp.shuffle(buffer_size=shuffle_size)
  # 1 read one json file
  self.dp = TextLineDataPipe(dp)
  self.dp = self.dp.prefetch(prefetch)