Update build dataset generator selection

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()