Fix positional argument error (#26)

* Fix positional argument error * Fix build dataset bug --------- Co-authored-by: Leng Yue <[email protected]>
fishaudio · Dec 20, 2023 · c6904b6 · c6904b6
1 parent 198ccfa
commit c6904b6
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tools/llama/build_dataset.py b/tools/llama/build_dataset.py
@@ -126,10 +126,14 @@ 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
+    generator_fn = (
+        task_generator_yaml(config)
+        if filelist is None
+        else task_generator_filelist(filelist)
+    )
 
     with Pool(num_worker) as p:
-        for result in tqdm(p.imap_unordered(run_task, generator_fn(config, filelist))):
+        for result in tqdm(p.imap_unordered(run_task, generator_fn)):
             dataset_fp.write(result)
 
     dataset_fp.close()