Zhiyu Wu Zhiyu Wu commited on
Commit
3d567c0
1 Parent(s): 81672d7

Fix: End of file handling in `benchmark.py` (#16)

Browse files

Co-authored-by: Zhiyu Wu <zhiyuwu@ampere01.eecs.umich.edu>

Files changed (1) hide show
  1. scripts/benchmark.py +3 -1
scripts/benchmark.py CHANGED
@@ -268,7 +268,7 @@ def main(
268
  model_path = model_path[:-1]
269
  model_name_cleaned = "--".join(model_path.split("/")[-2:])
270
  output_dir = f"{output_dir}/{task}/{model_name_cleaned}"
271
- output_csv_path = f"{output_dir}/benchmark.json"
272
  config_json_path = f"{output_dir}/config.json"
273
  table = Table(title="Benchmark")
274
  table.add_column("Configuration")
@@ -380,6 +380,8 @@ def main(
380
  convs.append(conv)
381
  if (len(convs) < batch): continue
382
  gen_params["prompt"] = prompts
 
 
383
 
384
  # Print input prompt.
385
  for i in range(len(convs)):
 
268
  model_path = model_path[:-1]
269
  model_name_cleaned = "--".join(model_path.split("/")[-2:])
270
  output_dir = f"{output_dir}/{task}/{model_name_cleaned}"
271
+ output_csv_path = f"{output_dir}/benchmark_batch_{batch}.json"
272
  config_json_path = f"{output_dir}/config.json"
273
  table = Table(title="Benchmark")
274
  table.add_column("Configuration")
 
380
  convs.append(conv)
381
  if (len(convs) < batch): continue
382
  gen_params["prompt"] = prompts
383
+ if end_of_file and len(prompts) == 0:
384
+ break
385
 
386
  # Print input prompt.
387
  for i in range(len(convs)):