浏览代码

Fix positional argument error (#26)

* Fix positional argument error

* Fix build dataset bug

---------

Co-authored-by: Leng Yue <lengyue@lengyue.me>
OedoSoldier 2 年之前
父节点
当前提交
c6904b6082
共有 1 个文件被更改,包括 6 次插入2 次删除
  1. 6 2
      tools/llama/build_dataset.py

+ 6 - 2
tools/llama/build_dataset.py

@@ -126,10 +126,14 @@ def run_task(task):
 @click.option("--num_worker", type=int, default=16)
 @click.option("--num_worker", type=int, default=16)
 def main(config, output, filelist, num_worker):
 def main(config, output, filelist, num_worker):
     dataset_fp = open(output, "wb")
     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:
     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.write(result)
 
 
     dataset_fp.close()
     dataset_fp.close()