Function main

inference.py:166–194

(config_path: str, task_path: str, output_path: str)

        print("======Write Json=======")

def main(config_path: str, task_path: str, output_path: str):
    os.makedirs("Check", exist_ok=True)
    os.makedirs("output", exist_ok=True)
    logger.info("started")
    yaml = ruamel.yaml.YAML()
    config=yaml.load(open(config_path, "r"))
    logger.info("config loaded")
    llm = build_llm(config["llm"]["type"], config["llm"]["args"])
    logger.info("llm built")
    system_prompt_template = config["system_prompt_template"]
    max_turns = config["max_turns"]
    test_data = json.load(open(task_path, "r"))
    logger.info(f"total tasks: {len(test_data)}")
    if os.path.exists(output_path):
        processed_ids = set(
            [json.loads(line)["index"] for line in open(output_path, "r")]
        )
    else:
        processed_ids = set()

    with ThreadPoolExecutor(max_workers=16) as executor:
        futures = [
            executor.submit(process_task, task, config, llm, system_prompt_template, max_turns, output_path, processed_ids)
            for task in test_data
        ]
        for future in futures:
            future.result()  # 等待所有任务完成

    logger.info("finished")

if __name__ == "__main__":
    fire.Fire(main)