From c6904b60822d875426c664e4f5c0771b3f42f9f3 Mon Sep 17 00:00:00 2001 From: OedoSoldier <31711261+OedoSoldier@users.noreply.github.com> Date: Wed, 20 Dec 2023 20:22:33 +0800 Subject: [PATCH] Fix positional argument error (#26) * Fix positional argument error * Fix build dataset bug --------- Co-authored-by: Leng Yue --- tools/llama/build_dataset.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tools/llama/build_dataset.py b/tools/llama/build_dataset.py index 822e7d48..666edb88 100644 --- 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()