Function process_single_file

CommonKG/create_kg.py:147–329

处理单个文件的三元组抽取

(corpus_path, task_conf, llm_processer, output_dir="output")

def process_single_file(corpus_path, task_conf, llm_processer, output_dir="output"):
    """处理单个文件的三元组抽取"""
    start_time = time.time()
    pedia_entity_path = task_conf["pedia_entity_path"]  # 头实体路径

    try:
    # 动态生成输出路径
        file_name = Path(corpus_path).stem
        output_subdir = Path(output_dir) / file_name
        if output_subdir.exists():
            logger.info(f"Target files: {file_name} already exists, overwrite\n")
        else:
            output_subdir.mkdir(parents=True, exist_ok=True)
        #移动chunk文件
        shutil.copy(corpus_path, output_subdir)
        
        # 初始化输出文件路径
        result_triple_path = output_subdir / f"new_triples_{file_name}.jsonl"
        next_layer_entities_path = output_subdir / f"next_layer_entities_{file_name}.txt"
        all_entities_path = output_subdir / f"all_entities_{file_name}.txt"
        match_words_path = output_subdir / f"match_words_{file_name}.jsonl"  ## 匹配结果路径


        # 当不跳过提取三元组时，进行多层次的实体匹配和三元组抽取
        if not task_conf["skip_extract_triple"]:

            # 加载到头实体路径进行处理，假设有第0层，那么头实体直接next_layer_entities_path来匹配
            head_entities = read_txt(pedia_entity_path)
            next_layer_entities = list(set([item.strip() for item in head_entities]))
            write_txt(next_layer_entities_path, next_layer_entities, mode="w")
            logger.info(f"Initialize next_layer_entities num: {len(next_layer_entities)}")
            
            # 初始化实体和三元组文件
            write_jsonl(data="", path=result_triple_path, mode="w")
            write_txt(data="", path=all_entities_path, mode="w")
            with open(corpus_path, "r") as f:
                corpusfiles=json.load(f)
            # 读取语料文件
            # corpusfiles = read_jsonl(corpus_path)
            logger.info(f"corpus paragraph num: {len(corpusfiles)}")

            for iter in range(task_conf["level_num"]):
                logger.info(f"Processing {file_name} | Iteration {iter+1}/{task_conf['level_num']}")
                layer_head_cnt, layer_tail_cnt, layer_triple_cnt = 0, 0, 0

                logger.info(f"[num_iteration]: {iter+1} ---------------------\n")
                logger.info("[corpus matching]-----------------------------------\n")

                # 检查文件是否存在，如果存在则删除
                if os.path.exists(match_words_path):
                    os.remove(match_words_path)
                            
                next_layer_entities = read_txt(next_layer_entities_path)
                next_layer_entities = [entity.strip("\n") for entity in next_layer_entities]

                source_id = "hash_code"
                text_key = "text"