Update build dataset generator selection

fishaudio · Dec 20, 2023 · a2e3f6e · a2e3f6e
1 parent 4f02d63
commit a2e3f6e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tools/llama/build_dataset.py b/tools/llama/build_dataset.py
@@ -126,10 +126,10 @@ def run_task(task):
 @click.option("--num_worker", type=int, default=16)
 def main(config, output, filelist, num_worker):
     dataset_fp = open(output, "wb")
+    generator_fn = task_generator_yaml if filelist is None else task_generator_filelist
+
     with Pool(num_worker) as p:
-        for result in tqdm(
-            p.imap_unordered(run_task, task_generator(config, filelist))
-        ):
+        for result in tqdm(p.imap_unordered(run_task, generator_fn(config, filelist))):
             dataset_fp.write(result)
 
     dataset_fp.close()