129 lines
5.0 KiB
Python
129 lines
5.0 KiB
Python
# 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.")
|