(args)
| 458 | |
| 459 | |
| 460 | def predict(args): |
| 461 | |
| 462 | |
| 463 | args.data_dir, args.split, data_subset = get_dataset_dir(args.data_name) |
| 464 | |
| 465 | print("Setting up his_index.json...") |
| 466 | make_his_idx(episode_dir, episodes_files) |
| 467 | |
| 468 | print("Merging Weights...") |
| 469 | merge_weight() |
| 470 | |
| 471 | print(f"Predicting on: {args.data_dir}/{args.split}") |
| 472 | print(f"Data subset: {data_subset}") |
| 473 | |
| 474 | if multiprocessing.get_start_method(allow_none=True) != "spawn": |
| 475 | multiprocessing.set_start_method("spawn", force=True) |
| 476 | |
| 477 | with ProcessPoolExecutor(max_workers=len(DEVICES),initializer=_init_llm,initargs=(args.model_path,)) as poolexec: |
| 478 | tasks = [] |
| 479 | print("Moving model to devices") |
| 480 | futures = [poolexec.submit(move_to, dev) for dev in DEVICES] |
| 481 | for fut in futures: |
| 482 | print(fut.result()) |
| 483 | |
| 484 | for dataset in data_subset: |
| 485 | save_dir = os.path.join(args.output_dir, dataset) |
| 486 | if not os.path.exists(save_dir): |
| 487 | os.makedirs(save_dir) |
| 488 | |
| 489 | episode_dir = os.path.join(args.data_dir, args.split, dataset) |
| 490 | output_file = os.path.join(save_dir, "predict.jsonl") |
| 491 | |
| 492 | # Get the list of all episodes files |
| 493 | if os.path.exists(episode_dir): |
| 494 | episodes_files = os.listdir(episode_dir) |
| 495 | else: |
| 496 | continue |
| 497 | |
| 498 | all_tasks = [] |
| 499 | print("Loading episodes") |
| 500 | |
| 501 | all_tasks = build_data_episodes( |
| 502 | episode_dir = episode_dir, |
| 503 | episodes_files = episodes_files, |
| 504 | dataset_name = args.data_name, |
| 505 | his_len = args.his_len |
| 506 | ) |
| 507 | |
| 508 | with open(output_file, "w", encoding="utf-8") as f_out: |
| 509 | print("Predicting") |
| 510 | tasks = [] |
| 511 | for task_value in all_tasks: |
| 512 | tasks.append(poolexec.submit(run_episode, **task_value)) |
| 513 | |
| 514 | for task in tqdm(as_completed(tasks), total=len(tasks), dynamic_ncols=True): |
| 515 | try: |
| 516 | episode = task.result() |
| 517 | episode_json = json.dumps(episode, ensure_ascii=False) |
no test coverage detected