| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502 |
- # # 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 "</think>" in content:
- # idx = content.index("</think>") + len("</think>")
- # 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 "</think>" in content:
- idx = content.index("</think>") + len("</think>")
- 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.")
|