# -*- coding: utf-8 -*-
|
# doc_split.py
|
# @author: lyg
|
# @date: 2025-04-24
|
# @version: 0.0.1
|
# @description: 文档处理，按页切割文档并提取每一页的文本，提文档的章节信息，将章节信息和页码信息关联起来，提取每一页的实体词。
|
|
from dataclasses import dataclass, field
|
import docx
|
import pymupdf
|
|
import pymupdf4llm
|
|
from knowledgebase.doc.entity_recognition import EntityRecognition
|
|
|
@dataclass
|
class ChapterInfo:
|
"""
|
表示章节信息的类，包含章节标题、起始页码、结束页码和层级。
|
"""
|
title: str # 章节标题
|
start_page: int # 起始页码
|
end_page: int # 结束页码
|
level: int # 章节层级
|
|
def __post_init__(self):
|
"""
|
在初始化后执行的钩子函数，用于验证输入参数的有效性。
|
"""
|
# 验证标题是否为空
|
if not self.title or not isinstance(self.title, str):
|
raise ValueError("标题必须是非空字符串")
|
|
# 验证页码是否为正整数
|
if not isinstance(self.start_page, int) or self.start_page < 1:
|
raise ValueError("起始页码必须是大于0的整数")
|
if not isinstance(self.end_page, int) or self.end_page < 1:
|
raise ValueError("结束页码必须是大于0的整数")
|
|
# 验证起始页码是否小于等于结束页码
|
if self.start_page > self.end_page:
|
raise ValueError("起始页码不能大于结束页码")
|
|
# 验证层级是否为正整数
|
if not isinstance(self.level, int) or self.level < 1:
|
raise ValueError("层级必须是大于0的整数")
|
|
|
@dataclass
|
class PageInfo:
|
"""
|
页面信息类，用于存储页面的基本信息。
|
|
参数：
|
- page_num (int): 页面编号，必须为正整数。
|
- chapter_info (str): 章节信息，描述当前页面所属章节。
|
- text (str, optional): 页面文本内容，默认为空字符串。
|
- entities (list[str]): 页面实体词列表。
|
|
异常：
|
- ValueError: 如果 `page_num` 不是正整数，或者 `chapter_info` 和 `text` 不是字符串。
|
"""
|
page_num: int
|
chapter_info: str
|
text: str = ""
|
entities: list[str] = field(default_factory=list)
|
|
def __post_init__(self):
|
"""
|
初始化后验证参数类型和值。
|
"""
|
if not isinstance(self.page_num, int) or self.page_num < 0:
|
raise ValueError("page_num 必须是正整数")
|
if not isinstance(self.chapter_info, str):
|
raise ValueError("chapter_info 必须是字符串")
|
if not isinstance(self.text, str):
|
raise ValueError("text 必须是字符串")
|
if self.entities is None or not isinstance(self.entities, list):
|
raise ValueError("entities 必须是列表")
|
|
|
class DocSplit:
|
"""
|
文档处理
|
按页切割文档并提取每一页都文本，使用
|
"""
|
pdf_file: str
|
page_infos: list[PageInfo]
|
chapter_infos: list[ChapterInfo]
|
|
def __init__(self, pdf_file):
|
self.pdf_file = pdf_file
|
self.page_infos = []
|
self.chapter_infos = []
|
self.doc = pymupdf.open(self.pdf_file)
|
self.extract_chapter_info()
|
self.extract_page_info()
|
|
def extract_chapter_info(self):
|
"""
|
提取章节信息
|
"""
|
toc = self.doc.get_toc()
|
for item in toc:
|
idx = toc.index(item)
|
end = len(self.doc)
|
for i in range(idx + 1, len(toc)):
|
if toc[i][0] >= item[0]:
|
end = toc[i][2]
|
break
|
self.chapter_infos.append(ChapterInfo(title=item[1], start_page=item[2], end_page=end, level=item[0]))
|
|
def extract_page_info(self):
|
"""
|
提取页面信息
|
"""
|
for page in self.doc:
|
# 通过pymupdf4llm获取页面文本，markdown格式
|
page_text = pymupdf4llm.to_markdown(self.doc, pages=[page.number])
|
# 创建页面信息对象
|
page_info = PageInfo(page_num=page.number, chapter_info="", text=page_text)
|
self.page_infos.append(page_info)
|
|
def get_page_info(self, page_num):
|
"""
|
获取指定页码的页面信息
|
"""
|
for page in self.page_infos:
|
if page.page_num == page_num:
|
return page
|
return None
|
|
|
# if __name__ == '__main__':
|
# ds = DocSplit("D:/workspace/PythonProjects/KnowledgeBase/doc/XA-5D无人机探测大纲（公开）111.pdf")
|
# print()
|
