Function main

CommonKG/create_kg.py:348–381

()

def main():
    ## 读取配置文件
    conf_path = "CommonKG/config/create_kg_conf_test.yaml"
    with open(conf_path, "r", encoding="utf-8") as file:
        args = yaml.safe_load(file)

    logger.info(f"args:\n{args}\n")    

    task_conf = args["task_conf"]  ## 任务参数

    # 迭代提取次数
    llm_conf = args["llm_conf"]  ## llm参数
    llm_processer = LLM_Processor(llm_conf)

    
    # 输入路径处理（支持文件/文件夹）
    input_path = task_conf["corpus_path"]
    output_dir = task_conf["output_dir"]
    if os.path.isfile(input_path):
        files_to_process = [input_path]
    elif os.path.isdir(input_path):
        files_to_process = [os.path.join(input_path, f) for f in os.listdir(input_path) if f.endswith(".jsonl")]
    else:
        raise ValueError(f"Invalid input path: {input_path}")

    # 批量处理
    success_count = 0
    
    # 移动chunk文件
    for corpus_path in tqdm(files_to_process, desc="Processing files"):
        if process_single_file(corpus_path, task_conf, llm_processer, output_dir):
            success_count += 1

    logger.info(f"Processed {success_count}/{len(files_to_process)} files successfully")