Просмотр исходного кода

Fix positional argument error (#26)

* Fix positional argument error

* Fix build dataset bug

---------

Co-authored-by: Leng Yue <lengyue@lengyue.me>
OedoSoldier 2 лет назад
Родитель
Сommit
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)
 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()