# # Copyright (c) Opendatalab. All rights reserved. # from loguru import logger # from openai import OpenAI # import json_repair # from mineru.backend.pipeline.pipeline_middle_json_mkcontent import merge_para_with_text # def llm_aided_title(page_info_list, title_aided_config): # client = OpenAI( # api_key=title_aided_config["api_key"], # base_url=title_aided_config["base_url"], # ) # title_dict = {} # origin_title_list = [] # i = 0 # for page_info in page_info_list: # blocks = page_info["para_blocks"] # for block in blocks: # if block["type"] == "title": # origin_title_list.append(block) # title_text = merge_para_with_text(block) # if 'line_avg_height' in block: # line_avg_height = block['line_avg_height'] # else: # title_block_line_height_list = [] # for line in block['lines']: # bbox = line['bbox'] # title_block_line_height_list.append(int(bbox[3] - bbox[1])) # if len(title_block_line_height_list) > 0: # line_avg_height = sum(title_block_line_height_list) / len(title_block_line_height_list) # else: # line_avg_height = int(block['bbox'][3] - block['bbox'][1]) # title_dict[f"{i}"] = [title_text, line_avg_height, int(page_info['page_idx']) + 1] # i += 1 # # logger.info(f"Title list: {title_dict}") # title_optimize_prompt = f"""输入的内容是一篇文档中所有标题组成的字典，请根据以下指南优化标题的结果，使结果符合正常文档的层次结构： # 1. 字典中每个value均为一个list，包含以下元素： # - 标题文本 # - 文本行高是标题所在块的平均行高 # - 标题所在的页码 # 2. 保留原始内容： # - 输入的字典中所有元素都是有效的，不能删除字典中的任何元素 # - 请务必保证输出的字典中元素的数量和输入的数量一致 # 3. 保持字典内key-value的对应关系不变 # 4. 优化层次结构： # - 根据标题内容的语义为每个标题元素添加适当的层次结构 # - 行高较大的标题一般是更高级别的标题 # - 标题从前至后的层级必须是连续的，不能跳过层级 # - 标题层级最多为4级，不要添加过多的层级 # - 优化后的标题只保留代表该标题的层级的整数，不要保留其他信息 # 5. 合理性检查与微调： # - 在完成初步分级后，仔细检查分级结果的合理性 # - 根据上下文关系和逻辑顺序，对不合理的分级进行微调 # - 确保最终的分级结果符合文档的实际结构和逻辑 # IMPORTANT: # 请直接返回优化过的由标题层级组成的字典，格式为{{标题id:标题层级}}，如下： # {{ # 0:1, # 1:2, # 2:2, # 3:3 # }} # 不需要对字典格式化，不需要返回任何其他信息。 # Input title list: # {title_dict} # Corrected title list: # """ # #5. # #- 字典中可能包含被误当成标题的正文，你可以通过将其层级标记为 0 来排除它们 # retry_count = 0 # max_retries = 3 # dict_completion = None # # Build API call parameters # api_params = { # "model": title_aided_config["model"], # "messages": [{'role': 'user', 'content': title_optimize_prompt}], # "temperature": 0.7, # "stream": True, # } # # Only add extra_body when explicitly specified in config # if "enable_thinking" in title_aided_config: # api_params["extra_body"] = {"enable_thinking": title_aided_config["enable_thinking"]} # while retry_count < max_retries: # try: # completion = client.chat.completions.create(**api_params) # content_pieces = [] # for chunk in completion: # if chunk.choices and chunk.choices[0].delta.content is not None: # content_pieces.append(chunk.choices[0].delta.content) # content = "".join(content_pieces).strip() # # logger.info(f"Title completion: {content}") # if "" in content: # idx = content.index("") + len("") # content = content[idx:].strip() # dict_completion = json_repair.loads(content) # dict_completion = {int(k): int(v) for k, v in dict_completion.items()} # # logger.info(f"len(dict_completion): {len(dict_completion)}, len(title_dict): {len(title_dict)}") # if len(dict_completion) == len(title_dict): # for i, origin_title_block in enumerate(origin_title_list): # origin_title_block["level"] = int(dict_completion[i]) # break # else: # logger.warning( # "The number of titles in the optimized result is not equal to the number of titles in the input.") # retry_count += 1 # except Exception as e: # logger.exception(e) # retry_count += 1 # if dict_completion is None: # logger.error("Failed to decode dict after maximum retries.") # Copyright (c) Opendatalab. All rights reserved. import re from loguru import logger from openai import AsyncOpenAI import json_repair import time from mineru.backend.pipeline.pipeline_middle_json_mkcontent import merge_para_with_text async def llm_aided_title(page_info_list, title_aided_config): client = AsyncOpenAI( api_key=title_aided_config["api_key"], base_url=title_aided_config["base_url"], ) title_dict = {} origin_title_list = [] i = 0 for page_info in page_info_list: blocks = page_info["para_blocks"] # 处理list类型的blocks，将其展开为text类型 processed_blocks = [] for block in blocks: if block.get("type") == "list": # 如果是list类型，提取其内部的blocks并添加到processed_blocks if "blocks" in block and isinstance(block["blocks"], list): processed_blocks.extend(block["blocks"]) else: # 非list类型直接添加 processed_blocks.append(block) page_info["para_blocks"] = processed_blocks # 更新blocks为处理后的列表 blocks = page_info["para_blocks"] for block in blocks: if block["type"] == "title": origin_title_list.append(block) title_text = merge_para_with_text(block) if 'line_avg_height' in block: line_avg_height = block['line_avg_height'] else: title_block_line_height_list = [] for line in block['lines']: bbox = line['bbox'] title_block_line_height_list.append(int(bbox[3] - bbox[1])) if len(title_block_line_height_list) > 0: line_avg_height = sum(title_block_line_height_list) / len(title_block_line_height_list) else: line_avg_height = int(block['bbox'][3] - block['bbox'][1]) # title_dict[f"{i}"] = [title_text, line_avg_height, int(page_info['page_idx']) + 1] title_dict[f"{i}"] = [title_text, int(page_info['page_idx']) + 1] i += 1 else: # 判断类型是否为text，并且内容是否以标题序号开头 if block["type"] == "text": title_text = merge_para_with_text(block) # 匹配标题序号格式：1、2、3、1.1、1.2、1.1.1、1.1.2等 # 支持：顿号、点号、空格作为分隔符，或者直接跟中文/字母 print(f'####文本: {title_text}') if re.match(r'^\d+(\.\d+)*([、.\s]|(?=[a-zA-Z\u4e00-\u9fa5]))', title_text): # block["type"] = "title" origin_title_list.append(block) print(f'####是否是标题: 是！') if 'line_avg_height' in block: line_avg_height = block['line_avg_height'] else: title_block_line_height_list = [] for line in block['lines']: bbox = line['bbox'] title_block_line_height_list.append(int(bbox[3] - bbox[1])) if len(title_block_line_height_list) > 0: line_avg_height = sum(title_block_line_height_list) / len(title_block_line_height_list) else: line_avg_height = int(block['bbox'][3] - block['bbox'][1]) # title_dict[f"{i}"] = [title_text, line_avg_height, int(page_info['page_idx']) + 1] title_dict[f"{i}"] = [title_text, int(page_info['page_idx']) + 1] i += 1 logger.info(f"Title list: {title_dict}") with open("data.txt", "w") as file: file.write(str(title_dict)) # time.sleep(10) # title_optimize_prompt = f"""输入的内容是一篇文档中所有标题组成的字典，请根据以下指南优化标题的结果，使结果符合正常文档的层次结构： # 1. 字典中每个value均为一个list，包含以下元素： # - 标题文本 # - 文本行高是标题所在块的平均行高 # - 标题所在的页码 # 2. 保留原始内容： # - 输入的字典中所有元素都是有效的，不能删除字典中的任何元素 # - 请务必保证输出的字典中元素的数量和输入的数量一致 # 3. 保持字典内key-value的对应关系不变 # 4. 优化层次结构： # - 为每个标题元素添加适当的层次结构 # - 行高较大的标题一般是更高级别的标题 # - 标题从前至后的层级必须是连续的，不能跳过层级 # - 标题层级最多为5级，不要添加过多的层级 # - 优化后的标题只保留代表该标题的层级的整数，不要保留其他信息 # - 一般的，“.”的个数不同的不是一个级别且“.”的个数越多的级别通常越低，例如“x.x通常比x.x.x级别高” # - 一般的，连续或相邻的x.x……等类似的是同级标题，例如“1.1.1、1.1.2、2.1.1、2.2.3或者1.1、2.1、3.1等类似的是同级别标题” # 5. 合理性检查与微调： # - 在完成初步分级后，仔细检查分级结果的合理性 # - 根据上下文关系和逻辑顺序，对不合理的分级进行微调 # - 确保最终的分级结果符合文档的实际结构和逻辑 # - 字典中可能包含被误当成标题的正文，你可以通过将其层级标记为 0 来排除它们 # - 一般的，如开头部分存在特殊字符，如*、#、&等，例如“1.0*DL、1.0#DL、1.0&DL”，你可以将这些字符的标题层级标记为 0 来排除它们 # 6. 生成树形路径： # - 从标题文本中提取编号（如"3.1 安全管理"提取"3.1"），无编号则用标题文本 # - 遇到1级标题时，以它作为新的根节点，路径就是它自身的编号/文本 # - 后续非1级标题的路径 = 当前根节点 + "->" + 各级父标题编号 + "->" + 自身编号 # - 遇到下一个1级标题时，切换为新的根节点，重复上述逻辑 # - 层级为0的非标题项，路径标记为"0" # 完整示例（注意paths的value是从标题文本提取的编号，不是标题id）： # 输入： # {{"0":["前言",24,1],"1":["目次",24,2],"2":["1 总则",22,3],"3":["1.1 为科学评价",18,3],"4":["1.2 本标准适用",18,3],"5":["3 检查评定项目",22,5],"6":["3.1 安全管理",20,5],"7":["3.1.3 保证项目",18,6],"8":["1 安全生产责任制",16,6],"9":["2 施工组织设计",16,6]}} # 输出： # {{"levels":{{"0":1,"1":1,"2":1,"3":2,"4":2,"5":1,"6":2,"7":3,"8":4,"9":4}},"paths":{{"0":"前言","1":"目次","2":"1","3":"1->1.1","4":"1->1.2","5":"3","6":"3->3.1","7":"3->3.1->3.1.3","8":"3->3.1->3.1.3->1","9":"3->3.1->3.1.3->2"}}}} # IMPORTANT: # 请返回一个JSON对象，包含两个字段： # - "levels": 层级字典，key是标题id，value是层级数字 # - "paths": 路径字典，key是标题id，value是从标题文本提取编号后构建的路径 # 不需要对JSON格式化，不需要返回任何其他信息。 # Input title list: # {title_dict} # Corrected title list: # """ title_optimize_prompt = f""" 输入的内容是一篇文档中所有标题组成的字典，请根据以下指南优化标题的结果，使结果符合下面规则的层次结构： 1. 字典中每个value均为一个list，包含以下元素： - 标题文本 - 标题所在的页码 2. 保留原始内容： - 输入的字典中所有元素都是有效的，不能删除字典中的任何元素 - 请务必保证输出的字典中元素的数量和输入的数量一致 3. 保持字典内key-value的对应关系不变 4. 优化层次结构（规则）： - 当输入标题的编号中包含“.”时，必须判定为“标题”且层级由编号形式严格决定，不得因结构合理性或正文语义而改变： - x “正文” 不含“.”的编号 → 必须为 2 级标题（如“1 总则”“2 术语”“3 基本规定”） - x.x “正文” → 必须为 3 级标题 - x.x.x “正文” → 必须为 4 级标题 - x.x.x.x “正文” → 必须为 5 级标题 即：层级 = 2 + 标题编号中“.”的数量 - 允许跳级，即x “正文”（2级标题）后面可以是x.x.x “正文”（4级标题） - 优化后的标题层级仅保留代表层级的整数，不要保留任何其他信息 - 对于仅包含简单序号形式的标题（如“1”“2”“3”），需要结合上下文进行判断： - 如果在一页中且连续，通常是上一个标题的子标题 此刻可将其视为子标题赋予层级，或标记为 0 - 若不在一页上，则很可能是2级标题 - 如果没有“.”，类似“一、”的视为2级标题，类似“（一）、”视为3级标题 5. 合理性检查与微调： - 在完成初步分级后，仔细检查分级结果是否符合规则，对规则中没有说明的标题按照合理的逻辑判定 - 不是标题的层级标记为0 - 一般的，如标题编号或文本中包含特殊字符（如 *、#、& 等）， 例如“1.0*DL “正文”、1.0#DL “正文”、1.0&DL “正文””，可将其视为非规范标题，并将其层级标记为 0 - 日期不是标题 6. 生成树形路径： - 从标题文本中提取编号（如"3.1 安全管理"提取"3.1"），无编号则用标题文本 - 遇到2级标题时，以它作为新的根节点，路径就是它自身的编号/文本 - 后续非2级标题的路径 = 当前根节点 + "->" + 各级父标题编号 + "->" + 自身编号，例如：“有标题3 xx、3.1 xx、3.1.1 xx、3.2 xx、3.2.1xx则输出：3、3->3.1、3->3.1->3.1.1、3->3.2、3->3.2->3.2.1” - 遇到下一个2级标题时，切换为新的根节点，重复上述逻辑 - 1级标题以及层级为0的非标题项，路径标记为"0" 完整示例（注意paths的value是从标题文本提取的编号，不是标题id）： 输入： {{"0":["前言",1],"1":["目次",2],"2":["1 xx",3],"3":["1.1 xx",3],"4":["1.1.1 xx",3],"5":["1.1.1.1 xx",3],"6":["1.2 xx",4],"7":["2 xx",5],"8":["2.1 xx",5],"9":["1 xx",6],"10":["2 xx",6],"11":["3 xx",6],"12":["2.1.1*DL xx",6],"13":["3 xx",7],"14":["3.1 xx",7],"15":["3.1.1 xx",7],"16":["3.1.2 xx",7]}} 输出： {{"levels":{{"0":1,"1":1,"2":2,"3":3,"4":4,"5":5,"6":3,"7":2,"8":3,"9":4,"10":4,"11":0,"12":0,"13":2,"14":3,"15":4,"16":4}},"paths":{{"0":"0","1":"0","2":"1","3":"1->1.1","4":"1->1.1->1.1.1","5":"1->1.1->1.1.1->1.1.1.1","6":"1->1.2","7":"2","8":"2->2.1","9":"2->2.1->1","10":"2->2.1->2","11":"2->2.1->3","12":"0","13":"3","14":"3->3.1","15":"3->3.1->3.1.1","16":"3->3.1->3.1.2"}}}} IMPORTANT: 请返回一个JSON对象，包含两个字段： - "levels": 层级字典，key是标题id，value是层级数字 - "paths": 路径字典，key是标题id，value是从标题文本提取编号后构建的路径 不需要对JSON格式化，不需要返回任何其他信息。 Input title list: {title_dict} Corrected title list: """ # - 字典中可能包含被误当成标题的正文，你可以通过将其层级标记为 0 来排除它们（符合“4. 优化层次结构（核心规则）”的除外） # - 编号结构是“x.x “正文”、x.x.x “正文”、…… “正文””的必须判定为标题 # - 禁止对已由“4. 优化层次结构（核心规则）”确定的标题层级进行升级或降级 # - 若其行高高于前后标题，可通过上述规则视为二级标题 # - 文本行高是标题所在块的平均行高（忽略） # - 若其前后标题的行高与该标题基本一致，且在结构上明显从属于前一个更高层级标题， # - 行高较大的标题通常优先视为更高层级标题，可作为辅助判断依据 # - 一般的，“.”的个数不同的标题不属于同一级别，且“.”越多层级通常越低 # 例如：“x.x”通常比“x.x.x”层级高 # - 一般的，连续或相邻的 x.x… 形式标题通常为同级标题， # 例如：“1.1.1、1.1.2、2.1.1、2.2.3” 或 “1.1、2.1、3.1” 等通常视为同一层级 # - 一般的，在非一级标题的标题后面出现“1、2、3、4”类似的，通常是子标题，例如“x.x……后面是1、2、……类似的那么1、2、……通常是x.x……的子标题” # - 标题从前至后的层级必须是连续的，不能出现层级跳跃（例如不能从2级直接跳到4级，有“.”的标题除外） # - 标题层级最多为5级，不允许超过该层级 # title_optimize_prompt = f""" # 输入的内容是一篇文档中所有标题组成的字典，请根据以下指南优化标题的结果，使结果符合正常文档的层次结构： # 1. 字典中每个value均为一个list，包含以下元素： # - 标题文本 # - 文本行高是标题所在块的平均行高 # - 标题所在的页码 # 2. 保留原始内容： # - 输入的字典中所有元素都是有效的，不能删除字典中的任何元素 # - 请务必保证输出的字典中元素的数量和输入的数量一致 # 3. 保持字典内key-value的对应关系不变 # 4. 优化层次结构（核心规则）： # - 仅当某一项被判定为“标题”时，才为其分配层级 # - 一般情况下，结构是"x.x.……"的是标题 # - 严格优先：标题编号中不包含“.”的（如“1 总则”“2 术语”“3 基本规定”），默认且强制视为二级标题（层级=2）。忽略任何上下文或行高干扰，直接应用此规则。 # - 标题编号中包含“.”的，其层级由“.”的个数决定： # - x.x → 3级标题 # - x.x.x → 4级标题 # - x.x.x.x → 5级标题 # 即：层级 = 2 + 标题编号中“.”的数量 # - 行高较大的标题可作为辅助判断依据，但不能覆盖默认规则；仅用于不确定简单序号时 # - 优化后的标题层级仅保留代表层级的整数，不要保留任何其他信息 # - 当你判定它为标题后且符合上面的描述（即“1 总则”“x.x”……）必须严格按照上面的要求进行定级，不得因上下文语义微调层级 # - 对于仅包含简单序号形式的标题（如“1”“2”“3”），需要结合上下文进行判断： # - 若其行高高于前后标题，且带完整文本（如“3 基本规定”），强制通过上述默认规则视为二级标题 # - 若其前后标题的行高与该标题基本一致，且在结构上明显从属于前一个更高层级标题，则可将其视为子标题，并根据上下文结构为其分配合理层级（例如：（一）、xx，1.xx 很明显1是（一）的子标题，或根据结构标记为 0） # - 若无法通过上下文和行高明确判断其为有效标题层级，或该类编号更可能是正文列表项而非标题，则应将其层级标记为 0 以排除 # - 如果没有“.”，一般的类似“一、”的为二级标题，类似“（一）、”为三级标题 # - 不存在连续的一级标题；当检测到结构、行高相同但语义和编号明显存在从属关系的标题时，即使其表征特征相同，也必须根据标题编号形式优先（默认二级），仅在编号形式不同时调整其后所有相关标题的层级结构，以确保文档层次符合常规规范结构。 # - 例如：一、xx，（一）、xx 可能分别是 二级标题和三级标题 # - 负面示例：勿将“3 基本规定”视为“2 术语”的子级，即使语义相关；它们是同级二级 # - 相同的结构标题级别相同，当你判定某一项为标题且赋予层级后，其它相同结构的也必须定义为标题并赋予相同层级 # - 例如：（一）、xx （二）、xx 同级，一、xx 二、xx 同级，1.0、2.0 同级，1.1、1.2同级 # - 处理目次干扰：在提取标题文本时，忽略附加页码（如“3 基本规定 6”提取为“3 基本规定”，不计“6”为编号部分） # 5. 合理性检查与微调： # - 在完成初步分级后，仔细检查分级结果的合理性，但微调仅限于将明显非标题标记为0，不得改变默认二级标题的层级 # - 根据上下文关系和逻辑顺序，对不合理的分级进行微调，但优先保持编号形式决定层级 # - 确保最终的分级结果符合文档的实际结构和逻辑 # - 字典中可能包含被误当成标题的正文，你可以通过将其层级标记为 0 来排除它们 # - 一般的，如标题编号或文本中包含特殊字符（如 *、#、& 等），例如“1.0*DL、1.0#DL、1.0&DL”，可将其视为非规范标题，并将其层级标记为 0 # 6. 生成树形路径： # - 从标题文本中提取编号（如"3.1 安全管理"提取"3.1"），无编号则用标题文本 # - 遇到1级标题时，以它作为新的根节点，路径就是它自身的编号/文本 # - 后续非1级标题的路径 = 当前根节点 + "->" + 各级父标题编号 + "->" + 自身编号，例如：“3->3.1->3.1.1” # - 遇到下一个1级标题时，切换为新的根节点，重复上述逻辑 # - 层级为0的非标题项，路径标记为"0" # - 相同层级的标题一定不是父子关系 # - 如果没有1级标题可以合理的根据2级标题进行上述编排，每个2级标题作为独立根（如"1"、"2"、"3"） # 完整示例（注意paths的value是从标题文本提取的编号，不是标题id）： # 输入： # {{"0":["前言",26,1],"1":["目次",26,2],"2":["1 总则",22,3],"3":["1.1 范围",18,3],"4":["1.1.1 术语定义",16,3],"5":["1.1.1.1 名词解释",14,3],"6":["1.2 规范性引用文件",18,4],"7":["2 技术要求",22,5],"8":["2.1 基本规定",18,5],"9":["1",16,6],"10":["2",16,6],"11":["3",12,6],"12":["2.1.1*DL 特殊说明",16,6],"13":["3 检查评定项目",22,7],"14":["3.1 安全管理",18,7],"15":["3.1.1 基本项目",16,7],"16":["3.1.2 一般项目",16,7]}} # 输出： # {{"levels":{{"0":1,"1":1,"2":2,"3":3,"4":4,"5":5,"6":3,"7":2,"8":3,"9":4,"10":4,"11":0,"12":0,"13":2,"14":3,"15":4,"16":4}},"paths":{{"0":"前言","1":"目次","2":"1","3":"1->1.1","4":"1->1.1->1.1.1","5":"1->1.1->1.1.1->1.1.1.1","6":"1->1.2","7":"2","8":"2->2.1","9":"2->2.1->1","10":"2->2.1->2","11":"0","12":"0","13":"3","14":"3->3.1","15":"3->3.1->3.1.1","16":"3->3.1->3.1.2"}}}} # IMPORTANT: # 请返回一个JSON对象，包含两个字段： # - "levels": 层级字典，key是标题id，value是层级数字 # - "paths": 路径字典，key是标题id，value是从标题文本提取编号后构建的路径 # 不需要对JSON格式化，不需要返回任何其他信息。 # Input title list: # {title_dict} # Corrected title list: # """ retry_count = 0 max_retries = 3 dict_completion = None while retry_count < max_retries: try: completion = await client.chat.completions.create( model=title_aided_config["model"], messages=[ {'role': 'user', 'content': title_optimize_prompt}], temperature=0.1, stream=True, max_tokens=16384, # + top_p=0.95, # + extra_body={ "top_k": 1, # 扩展：top-k采样 "min_p": 0.0, # 扩展：min-p采样 } ) content_pieces = [] async for chunk in completion: if chunk.choices and chunk.choices[0].delta.content is not None: content_pieces.append(chunk.choices[0].delta.content) logger.info(chunk) content = "".join(content_pieces).strip() with open("data_title.txt", "w") as file: file.write(str(content)) # logger.info(f"Title completion: {content}") # time.sleep(10) if "" in content: idx = content.index("") + len("") content = content[idx:].strip() # 解析嵌套JSON对象 result = json_repair.loads(content) levels_dict = result.get("levels", {}) paths_dict = result.get("paths", {}) dict_completion = {int(k): int(v) for k, v in levels_dict.items()} path_dict = {int(k): str(v) for k, v in paths_dict.items()} # logger.info(f"len(dict_completion): {len(dict_completion)}, len(title_dict): {len(title_dict)}") if len(dict_completion) == len(title_dict): for i, origin_title_block in enumerate(origin_title_list): # origin_title_block["level"] = int(dict_completion[i]) level = int(dict_completion[i]) origin_title_block["level"] = level # 添加树形路径 if i in path_dict: origin_title_block["title_path"] = path_dict[i] # 如果原本是text类型但level>0,说明LLM判断它是标题 if origin_title_block["type"] == "text" and level > 0: origin_title_block["type"] = "title" # 如果原本是title类型但level==0,说明LLM判断它不是标题 elif origin_title_block["type"] == "title" and level == 0: origin_title_block["type"] = "text" # logger.info(f"origin_title_block: {origin_title_list}") break else: logger.warning( "The number of titles in the optimized result is not equal to the number of titles in the input.") retry_count += 1 except Exception as e: logger.exception(e) retry_count += 1 if dict_completion is None: logger.error("Failed to decode dict after maximum retries.")