magic-pdf 0.10.6__py3-none-any.whl → 1.0.0__py3-none-any.whl

magic_pdf/model/sub_modules/table/rapidtable/rapid_table.py

@@ -1,16 +1,51 @@
+import cv2
 import numpy as np
+import torch
+from loguru import logger
 from rapid_table import RapidTable
-from rapidocr_paddle import RapidOCR
 
 
 class RapidTableModel(object):
-    def __init__(self):
+    def __init__(self, ocr_engine):
         self.table_model = RapidTable()
-        self.ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
+        # if ocr_engine is None:
+        #     self.ocr_model_name = "RapidOCR"
+        #     if torch.cuda.is_available():
+        #         from rapidocr_paddle import RapidOCR
+        #         self.ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
+        #     else:
+        #         from rapidocr_onnxruntime import RapidOCR
+        #         self.ocr_engine = RapidOCR()
+        # else:
+        #     self.ocr_model_name = "PaddleOCR"
+        #     self.ocr_engine = ocr_engine
+
+        self.ocr_model_name = "RapidOCR"
+        if torch.cuda.is_available():
+            from rapidocr_paddle import RapidOCR
+            self.ocr_engine = RapidOCR(det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True)
+        else:
+            from rapidocr_onnxruntime import RapidOCR
+            self.ocr_engine = RapidOCR()
 
     def predict(self, image):
-        ocr_result, _ = self.ocr_engine(np.asarray(image))
-        if ocr_result is None:
+
+        if self.ocr_model_name == "RapidOCR":
+            ocr_result, _ = self.ocr_engine(np.asarray(image))
+        elif self.ocr_model_name == "PaddleOCR":
+            bgr_image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
+            ocr_result = self.ocr_engine.ocr(bgr_image)[0]
+            if ocr_result:
+                ocr_result = [[item[0], item[1][0], item[1][1]] for item in ocr_result if
+                          len(item) == 2 and isinstance(item[1], tuple)]
+            else:
+                ocr_result = None
+        else:
+            logger.error("OCR model not supported")
+            ocr_result = None
+
+        if ocr_result:
+            html_code, table_cell_bboxes, elapse = self.table_model(np.asarray(image), ocr_result)
+            return html_code, table_cell_bboxes, elapse
+        else:
             return None, None, None
-        html_code, table_cell_bboxes, elapse = self.table_model(np.asarray(image), ocr_result)
-        return html_code, table_cell_bboxes, elapse

magic_pdf/operators/__init__.py

@@ -0,0 +1,94 @@
+from abc import ABC, abstractmethod
+from typing import Callable
+
+from magic_pdf.data.data_reader_writer import DataWriter
+from magic_pdf.data.dataset import Dataset
+from magic_pdf.operators.pipes import PipeResult
+
+
+class InferenceResultBase(ABC):
+
+    @abstractmethod
+    def __init__(self, inference_results: list, dataset: Dataset):
+        """Initialized method.
+
+        Args:
+            inference_results (list): the inference result generated by model
+            dataset (Dataset): the dataset related with model inference result
+        """
+        pass
+
+    @abstractmethod
+    def draw_model(self, file_path: str) -> None:
+        """Draw model inference result.
+
+        Args:
+            file_path (str): the output file path
+        """
+        pass
+
+    @abstractmethod
+    def dump_model(self, writer: DataWriter, file_path: str):
+        """Dump model inference result to file.
+
+        Args:
+            writer (DataWriter): writer handle
+            file_path (str): the location of target file
+        """
+        pass
+
+    @abstractmethod
+    def get_infer_res(self):
+        """Get the inference result.
+
+        Returns:
+            list: the inference result generated by model
+        """
+        pass
+
+    @abstractmethod
+    def apply(self, proc: Callable, *args, **kwargs):
+        """Apply callable method which.
+
+        Args:
+            proc (Callable): invoke proc as follows:
+                proc(inference_result, *args, **kwargs)
+
+        Returns:
+            Any: return the result generated by proc
+        """
+        pass
+
+    def pipe_txt_mode(
+        self,
+        imageWriter: DataWriter,
+        start_page_id=0,
+        end_page_id=None,
+        debug_mode=False,
+        lang=None,
+    ) -> PipeResult:
+        """Post-proc the model inference result, Extract the text using the
+        third library, such as `pymupdf`
+
+        Args:
+            imageWriter (DataWriter): the image writer handle
+            start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process
+            end_page_id (int, optional):  Defaults to the last page index of dataset. Let user select some pages He/She want to process
+            debug_mode (bool, optional): Defaults to False. will dump more log if enabled
+            lang (str, optional): Defaults to None.
+
+        Returns:
+            PipeResult: the result
+        """
+        pass
+
+    @abstractmethod
+    def pipe_ocr_mode(
+        self,
+        imageWriter: DataWriter,
+        start_page_id=0,
+        end_page_id=None,
+        debug_mode=False,
+        lang=None,
+    ) -> PipeResult:
+        pass

magic_pdf/{model/operators.py → operators/models.py}

@@ -7,13 +7,11 @@ from magic_pdf.config.constants import PARSE_TYPE_OCR, PARSE_TYPE_TXT
 from magic_pdf.config.enums import SupportedPdfParseMethod
 from magic_pdf.data.data_reader_writer import DataWriter
 from magic_pdf.data.dataset import Dataset
-from magic_pdf.filter import classify
 from magic_pdf.libs.draw_bbox import draw_model_bbox
 from magic_pdf.libs.version import __version__
-from magic_pdf.model import InferenceResultBase
+from magic_pdf.operators.pipes import PipeResult
 from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union
-from magic_pdf.pipe.operators import PipeResult
-
+from magic_pdf.operators import InferenceResultBase
 
 class InferenceResult(InferenceResultBase):
     def __init__(self, inference_results: list, dataset: Dataset):
@@ -71,40 +69,6 @@ class InferenceResult(InferenceResultBase):
         """
         return proc(copy.deepcopy(self._infer_res), *args, **kwargs)
 
-    def pipe_auto_mode(
-        self,
-        imageWriter: DataWriter,
-        start_page_id=0,
-        end_page_id=None,
-        debug_mode=False,
-        lang=None,
-    ) -> PipeResult:
-        """Post-proc the model inference result.
-            step1: classify the dataset type
-            step2: based the result of step1, using `pipe_txt_mode` or `pipe_ocr_mode`
-
-        Args:
-            imageWriter (DataWriter): the image writer handle
-            start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process
-            end_page_id (int, optional):  Defaults to the last page index of dataset. Let user select some pages He/She want to process
-            debug_mode (bool, optional): Defaults to False. will dump more log if enabled
-            lang (str, optional): Defaults to None.
-
-        Returns:
-            PipeResult: the result
-        """
-
-        pdf_proc_method = classify(self._dataset.data_bits())
-
-        if pdf_proc_method == SupportedPdfParseMethod.TXT:
-            return self.pipe_txt_mode(
-                imageWriter, start_page_id, end_page_id, debug_mode, lang
-            )
-        else:
-            return self.pipe_ocr_mode(
-                imageWriter, start_page_id, end_page_id, debug_mode, lang
-            )
-
     def pipe_txt_mode(
         self,
         imageWriter: DataWriter,

magic_pdf/{pipe/operators.py → operators/pipes.py}

@@ -1,7 +1,7 @@
+import copy
 import json
 import os
 from typing import Callable
-import copy
 
 from magic_pdf.config.make_content_config import DropMode, MakeMode
 from magic_pdf.data.data_reader_writer import DataWriter
@@ -23,12 +23,34 @@ class PipeResult:
         self._pipe_res = pipe_res
         self._dataset = dataset
 
+    def get_markdown(
+        self,
+        img_dir_or_bucket_prefix: str,
+        drop_mode=DropMode.NONE,
+        md_make_mode=MakeMode.MM_MD,
+    ) -> str:
+        """Get markdown content.
+
+        Args:
+            img_dir_or_bucket_prefix (str): The s3 bucket prefix or local file directory which used to store the figure
+            drop_mode (str, optional): Drop strategy when some page which is corrupted or inappropriate. Defaults to DropMode.NONE.
+            md_make_mode (str, optional): The content Type of Markdown be made. Defaults to MakeMode.MM_MD.
+
+        Returns:
+            str: return markdown content
+        """
+        pdf_info_list = self._pipe_res['pdf_info']
+        md_content = union_make(
+            pdf_info_list, md_make_mode, drop_mode, img_dir_or_bucket_prefix
+        )
+        return md_content
+
     def dump_md(
         self,
         writer: DataWriter,
         file_path: str,
         img_dir_or_bucket_prefix: str,
-        drop_mode=DropMode.WHOLE_PDF,
+        drop_mode=DropMode.NONE,
         md_make_mode=MakeMode.MM_MD,
     ):
         """Dump The Markdown.
@@ -37,36 +59,68 @@ class PipeResult:
             writer (DataWriter): File writer handle
             file_path (str): The file location of markdown
             img_dir_or_bucket_prefix (str): The s3 bucket prefix or local file directory which used to store the figure
-            drop_mode (str, optional): Drop strategy when some page which is corrupted or inappropriate. Defaults to DropMode.WHOLE_PDF.
+            drop_mode (str, optional): Drop strategy when some page which is corrupted or inappropriate. Defaults to DropMode.NONE.
             md_make_mode (str, optional): The content Type of Markdown be made. Defaults to MakeMode.MM_MD.
         """
-        pdf_info_list = self._pipe_res['pdf_info']
-        md_content = union_make(
-            pdf_info_list, md_make_mode, drop_mode, img_dir_or_bucket_prefix
+
+        md_content = self.get_markdown(
+            img_dir_or_bucket_prefix, drop_mode=drop_mode, md_make_mode=md_make_mode
         )
         writer.write_string(file_path, md_content)
 
-    def dump_content_list(
-        self, writer: DataWriter, file_path: str, image_dir_or_bucket_prefix: str
-    ):
-        """Dump Content List.
+    def get_content_list(
+        self,
+        image_dir_or_bucket_prefix: str,
+        drop_mode=DropMode.NONE,
+    ) -> str:
+        """Get Content List.
 
         Args:
-            writer (DataWriter): File writer handle
-            file_path (str): The file location of content list
             image_dir_or_bucket_prefix (str): The s3 bucket prefix or local file directory which used to store the figure
+            drop_mode (str, optional): Drop strategy when some page which is corrupted or inappropriate. Defaults to DropMode.NONE.
+
+        Returns:
+            str: content list content
         """
         pdf_info_list = self._pipe_res['pdf_info']
         content_list = union_make(
             pdf_info_list,
             MakeMode.STANDARD_FORMAT,
-            DropMode.NONE,
+            drop_mode,
             image_dir_or_bucket_prefix,
         )
+        return content_list
+
+    def dump_content_list(
+        self,
+        writer: DataWriter,
+        file_path: str,
+        image_dir_or_bucket_prefix: str,
+        drop_mode=DropMode.NONE,
+    ):
+        """Dump Content List.
+
+        Args:
+            writer (DataWriter): File writer handle
+            file_path (str): The file location of content list
+            image_dir_or_bucket_prefix (str): The s3 bucket prefix or local file directory which used to store the figure
+            drop_mode (str, optional): Drop strategy when some page which is corrupted or inappropriate. Defaults to DropMode.NONE.
+        """
+        content_list = self.get_content_list(
+            image_dir_or_bucket_prefix, drop_mode=drop_mode,
+        )
         writer.write_string(
             file_path, json.dumps(content_list, ensure_ascii=False, indent=4)
         )
 
+    def get_middle_json(self) -> str:
+        """Get middle json.
+
+        Returns:
+            str: The content of middle json
+        """
+        return json.dumps(self._pipe_res, ensure_ascii=False, indent=4)
+
     def dump_middle_json(self, writer: DataWriter, file_path: str):
         """Dump the result of pipeline.
 
@@ -74,9 +128,8 @@ class PipeResult:
             writer (DataWriter): File writer handler
             file_path (str): The file location of middle json
         """
-        writer.write_string(
-            file_path, json.dumps(self._pipe_res, ensure_ascii=False, indent=4)
-        )
+        middle_json = self.get_middle_json()
+        writer.write_string(file_path, middle_json)
 
     def draw_layout(self, file_path: str) -> None:
         """Draw the layout.
@@ -123,7 +176,7 @@ class PipeResult:
         Returns:
             str: compress the pipeline result and return
         """
-        return JsonCompressor.compress_json(self.pdf_mid_data)
+        return JsonCompressor.compress_json(self._pipe_res)
 
     def apply(self, proc: Callable, *args, **kwargs):
         """Apply callable method which.

magic_pdf/pdf_parse_union_core_v2.py

@@ -1,5 +1,6 @@
 import copy
 import os
+import re
 import statistics
 import time
 from typing import List
@@ -13,11 +14,12 @@ from magic_pdf.config.ocr_content_type import BlockType, ContentType
 from magic_pdf.data.dataset import Dataset, PageableData
 from magic_pdf.libs.boxbase import calculate_overlap_area_in_bbox1_area_ratio
 from magic_pdf.libs.clean_memory import clean_memory
-from magic_pdf.libs.config_reader import get_local_layoutreader_model_dir
+from magic_pdf.libs.config_reader import get_local_layoutreader_model_dir, get_llm_aided_config, get_device
 from magic_pdf.libs.convert_utils import dict_to_list
 from magic_pdf.libs.hash_utils import compute_md5
 from magic_pdf.libs.pdf_image_tools import cut_image_to_pil_image
 from magic_pdf.model.magic_model import MagicModel
+from magic_pdf.post_proc.llm_aided import llm_aided_formula, llm_aided_text, llm_aided_title
 
 try:
     import torchtext
@@ -28,15 +30,15 @@ except ImportError:
     pass
 
 from magic_pdf.model.sub_modules.model_init import AtomModelSingleton
-from magic_pdf.para.para_split_v3 import para_split
+from magic_pdf.post_proc.para_split_v3 import para_split
 from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component_v2
 from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table
 from magic_pdf.pre_proc.ocr_detect_all_bboxes import ocr_prepare_bboxes_for_layout_split_v2
 from magic_pdf.pre_proc.ocr_dict_merge import fill_spans_in_blocks, fix_block_spans_v2, fix_discarded_block
-from magic_pdf.pre_proc.ocr_span_list_modify import get_qa_need_list_v2, remove_overlaps_low_confidence_spans, remove_overlaps_min_spans
+from magic_pdf.pre_proc.ocr_span_list_modify import get_qa_need_list_v2, remove_overlaps_low_confidence_spans, \
+    remove_overlaps_min_spans, check_chars_is_overlap_in_span
 
 os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1'  # 禁止albumentations检查更新
-os.environ['YOLO_VERBOSE'] = 'False'  # disable yolo logger
 
 
 def __replace_STX_ETX(text_str: str):
@@ -63,11 +65,22 @@ def __replace_0xfffd(text_str: str):
         return s
     return text_str
 
+
+# 连写字符拆分
+def __replace_ligatures(text: str):
+    ligatures = {
+        'fi': 'fi', 'fl': 'fl', 'ff': 'ff', 'ffi': 'ffi', 'ffl': 'ffl', 'ſt': 'ft', 'st': 'st'
+    }
+    return re.sub('|'.join(map(re.escape, ligatures.keys())), lambda m: ligatures[m.group()], text)
+
+
 def chars_to_content(span):
     # 检查span中的char是否为空
     if len(span['chars']) == 0:
         pass
         # span['content'] = ''
+    elif check_chars_is_overlap_in_span(span['chars']):
+        pass
     else:
         # 先给chars按char['bbox']的中心点的x坐标排序
         span['chars'] = sorted(span['chars'], key=lambda x: (x['bbox'][0] + x['bbox'][2]) / 2)
@@ -78,11 +91,16 @@ def chars_to_content(span):
 
         content = ''
         for char in span['chars']:
-            # 如果下一个char的x0和上一个char的x1距离超过一个字符宽度，则需要在中间插入一个空格
-            if char['bbox'][0] - span['chars'][span['chars'].index(char) - 1]['bbox'][2] > char_avg_width:
-                content += ' '
-            content += char['c']
 
+            # 如果下一个char的x0和上一个char的x1距离超过0.25个字符宽度，则需要在中间插入一个空格
+            char1 = char
+            char2 = span['chars'][span['chars'].index(char) + 1] if span['chars'].index(char) + 1 < len(span['chars']) else None
+            if char2 and char2['bbox'][0] - char1['bbox'][2] > char_avg_width * 0.25 and char['c'] != ' ' and char2['c'] != ' ':
+                content += f"{char['c']} "
+            else:
+                content += char['c']
+
+        content = __replace_ligatures(content)
         span['content'] = __replace_0xfffd(content)
 
     del span['chars']
@@ -98,6 +116,10 @@ def fill_char_in_spans(spans, all_chars):
     spans = sorted(spans, key=lambda x: x['bbox'][1])
 
     for char in all_chars:
+        # 跳过非法bbox的char
+        x1, y1, x2, y2 = char['bbox']
+        if abs(x1 - x2) <= 0.01 or abs(y1 - y2) <= 0.01:
+            continue
         for span in spans:
             if calculate_char_in_span(char['bbox'], span['bbox'], char['c']):
                 span['chars'].append(char)
@@ -152,14 +174,16 @@ def calculate_char_in_span(char_bbox, span_bbox, char, span_height_radio=0.33):
 
 
 def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang):
+    # cid用0xfffd表示，连字符拆开
+    # text_blocks_raw = pdf_page.get_text('rawdict', flags=fitz.TEXT_PRESERVE_WHITESPACE | fitz.TEXT_MEDIABOX_CLIP)['blocks']
 
-    text_blocks_raw = pdf_page.get_text('rawdict', flags=fitz.TEXT_PRESERVE_WHITESPACE | fitz.TEXT_MEDIABOX_CLIP)['blocks']
-
+    # cid用0xfffd表示，连字符不拆开
+    text_blocks_raw = pdf_page.get_text('rawdict', flags=fitz.TEXT_PRESERVE_LIGATURES | fitz.TEXT_PRESERVE_WHITESPACE | fitz.TEXT_MEDIABOX_CLIP)['blocks']
     all_pymu_chars = []
     for block in text_blocks_raw:
         for line in block['lines']:
             cosine, sine = line['dir']
-            if abs (cosine) < 0.9 or abs(sine) > 0.1:
+            if abs(cosine) < 0.9 or abs(sine) > 0.1:
                 continue
             for span in line['spans']:
                 all_pymu_chars.extend(span['chars'])
@@ -255,19 +279,23 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang
     return spans
 
 
-def replace_text_span(pymu_spans, ocr_spans):
-    return list(filter(lambda x: x['type'] != ContentType.Text, ocr_spans)) + pymu_spans
-
-
 def model_init(model_name: str):
     from transformers import LayoutLMv3ForTokenClassification
-
+    device = get_device()
     if torch.cuda.is_available():
         device = torch.device('cuda')
         if torch.cuda.is_bf16_supported():
             supports_bfloat16 = True
         else:
             supports_bfloat16 = False
+    elif str(device).startswith("npu"):
+        import torch_npu
+        if torch_npu.npu.is_available():
+            device = torch.device('npu')
+            supports_bfloat16 = False
+        else:
+            device = torch.device('cpu')
+            supports_bfloat16 = False
     else:
         device = torch.device('cpu')
         supports_bfloat16 = False
@@ -345,6 +373,8 @@ def cal_block_index(fix_blocks, sorted_bboxes):
         # 使用xycut排序
         block_bboxes = []
         for block in fix_blocks:
+            # 如果block['bbox']任意值小于0，将其置为0
+            block['bbox'] = [max(0, x) for x in block['bbox']]
             block_bboxes.append(block['bbox'])
 
             # 删除图表body block中的虚拟line信息, 并用real_lines信息回填
@@ -738,6 +768,11 @@ def parse_page_core(
     """重排block"""
     sorted_blocks = sorted(fix_blocks, key=lambda b: b['index'])
 
+    """block内重排(img和table的block内多个caption或footnote的排序)"""
+    for block in sorted_blocks:
+        if block['type'] in [BlockType.Image, BlockType.Table]:
+            block['blocks'] = sorted(block['blocks'], key=lambda b: b['index'])
+
     """获取QA需要外置的list"""
     images, tables, interline_equations = get_qa_need_list_v2(sorted_blocks)
 
@@ -819,13 +854,29 @@ def pdf_parse_union(
     """分段"""
     para_split(pdf_info_dict)
 
+    """llm优化"""
+    llm_aided_config = get_llm_aided_config()
+    if llm_aided_config is not None:
+        """公式优化"""
+        formula_aided_config = llm_aided_config.get('formula_aided', None)
+        if formula_aided_config is not None:
+            llm_aided_formula(pdf_info_dict, formula_aided_config)
+        """文本优化"""
+        text_aided_config = llm_aided_config.get('text_aided', None)
+        if text_aided_config is not None:
+            llm_aided_text(pdf_info_dict, text_aided_config)
+        """标题优化"""
+        title_aided_config = llm_aided_config.get('title_aided', None)
+        if title_aided_config is not None:
+            llm_aided_title(pdf_info_dict, title_aided_config)
+
     """dict转list"""
     pdf_info_list = dict_to_list(pdf_info_dict)
     new_pdf_info_dict = {
         'pdf_info': pdf_info_list,
     }
 
-    clean_memory()
+    clean_memory(get_device())
 
     return new_pdf_info_dict
 

magic_pdf/post_proc/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Opendatalab. All rights reserved.

magic_pdf/post_proc/llm_aided.py

@@ -0,0 +1,133 @@
+# Copyright (c) Opendatalab. All rights reserved.
+import json
+from loguru import logger
+from magic_pdf.dict2md.ocr_mkcontent import merge_para_with_text
+from openai import OpenAI
+
+
+#@todo: 有的公式以"\"结尾，这样会导致尾部拼接的"$"被转义，也需要修复
+formula_optimize_prompt = """请根据以下指南修正LaTeX公式的错误，确保公式能够渲染且符合原始内容：
+
+1. 修正渲染或编译错误：
+    - Some syntax errors such as mismatched/missing/extra tokens. Your task is to fix these syntax errors and make sure corrected results conform to latex math syntax principles.
+    - 包含KaTeX不支持的关键词等原因导致的无法编译或渲染的错误
+
+2. 保留原始信息：
+   - 保留原始公式中的所有重要信息
+   - 不要添加任何原始公式中没有的新信息
+
+IMPORTANT:请仅返回修正后的公式，不要包含任何介绍、解释或元数据。
+
+LaTeX recognition result:
+$FORMULA
+
+Your corrected result:
+"""
+
+text_optimize_prompt = f"""请根据以下指南修正OCR引起的错误，确保文本连贯并符合原始内容：
+
+1. 修正OCR引起的拼写错误和错误：
+   - 修正常见的OCR错误（例如，'rn' 被误读为 'm'）
+   - 使用上下文和常识进行修正
+   - 只修正明显的错误，不要不必要的修改内容
+   - 不要添加额外的句号或其他不必要的标点符号
+
+2. 保持原始结构：
+   - 保留所有标题和子标题
+
+3. 保留原始内容：
+   - 保留原始文本中的所有重要信息
+   - 不要添加任何原始文本中没有的新信息
+   - 保留段落之间的换行符
+
+4. 保持连贯性：
+   - 确保内容与前文顺畅连接
+   - 适当处理在句子中间开始或结束的文本
+   
+5. 修正行内公式：
+   - 去除行内公式前后多余的空格
+   - 修正公式中的OCR错误
+   - 确保公式能够通过KaTeX渲染
+   
+6. 修正全角字符
+    - 修正全角标点符号为半角标点符号
+    - 修正全角字母为半角字母
+    - 修正全角数字为半角数字
+
+IMPORTANT:请仅返回修正后的文本，保留所有原始格式，包括换行符。不要包含任何介绍、解释或元数据。
+
+Previous context:
+
+Current chunk to process:
+
+Corrected text:
+"""
+
+def llm_aided_formula(pdf_info_dict, formula_aided_config):
+    pass
+
+def llm_aided_text(pdf_info_dict, text_aided_config):
+    pass
+
+def llm_aided_title(pdf_info_dict, 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_num, page in pdf_info_dict.items():
+        blocks = page["para_blocks"]
+        for block in blocks:
+            if block["type"] == "title":
+                origin_title_list.append(block)
+                title_text = merge_para_with_text(block)
+                title_dict[f"{i}"] = title_text
+                i += 1
+    # logger.info(f"Title list: {title_dict}")
+
+    title_optimize_prompt = f"""输入的内容是一篇文档中所有标题组成的字典，请根据以下指南优化标题的结果，使结果符合正常文档的层次结构：
+
+1. 保留原始内容：
+    - 输入的字典中所有元素都是有效的，不能删除字典中的任何元素
+    - 请务必保证输出的字典中元素的数量和输入的数量一致
+
+2. 保持字典内key-value的对应关系不变
+
+3. 优化层次结构：
+    - 为每个标题元素添加适当的层次结构
+    - 标题层级应具有连续性，不能跳过某一层级
+    - 标题层级最多为4级，不要添加过多的层级
+    - 优化后的标题为一个整数，代表该标题的层级
+
+IMPORTANT: 
+请直接返回优化过的由标题层级组成的json，返回的json不需要格式化。
+
+Input title list:
+{title_dict}
+
+Corrected title list:
+"""
+
+    completion = client.chat.completions.create(
+        model=title_aided_config["model"],
+        messages=[
+            {'role': 'user', 'content': title_optimize_prompt}],
+        temperature=0.7,
+    )
+
+    json_completion = json.loads(completion.choices[0].message.content)
+
+    # logger.info(f"Title completion: {json_completion}")
+
+    # logger.info(f"len(json_completion): {len(json_completion)}, len(title_dict): {len(title_dict)}")
+    if len(json_completion) == len(title_dict):
+        try:
+            for i, origin_title_block in enumerate(origin_title_list):
+               origin_title_block["level"] = int(json_completion[str(i)])
+        except Exception as e:
+            logger.exception(e)
+    else:
+        logger.error("The number of titles in the optimized result is not equal to the number of titles in the input.")
+