"""PDF Knowledge.""" import json import os import re from collections import defaultdict from typing import Any, Dict, List, Optional, Union from dbgpt.component import logger from dbgpt.core import Document from dbgpt.rag.knowledge.base import ( ChunkStrategy, DocumentType, Knowledge, KnowledgeType, ) class PDFKnowledge(Knowledge): """PDF Knowledge.""" def __init__( self, file_path: Optional[str] = None, knowledge_type: KnowledgeType = KnowledgeType.DOCUMENT, loader: Optional[Any] = None, language: Optional[str] = "zh", metadata: Optional[Dict[str, Union[str, List[str]]]] = None, **kwargs: Any, ) -> None: """Create PDF Knowledge with Knowledge arguments. Args: file_path(str, optional): file path knowledge_type(KnowledgeType, optional): knowledge type loader(Any, optional): loader language(str, optional): language """ super().__init__( path=file_path, knowledge_type=knowledge_type, data_loader=loader, metadata=metadata, **kwargs, ) self._language = language self._pdf_processor = PDFProcessor(filepath=self._path) self.all_title: List[dict] = [] self.all_text: List[dict] = [] def process_text_data(self): """Text data processing to level 1 and level 2 titles.""" for i, data in enumerate(self.all_text): # data = self.all_text[i] inside_content = data.get("inside") content_type = data.get("type") if content_type == "text": # use regex to match the first level title first_level_match = re.match( r"§(\d+)+([\u4e00-\u9fa5]+)", inside_content.strip() ) second_level_match = re.match( r"(\d+\.\d+)([\u4e00-\u9fa5]+)", inside_content.strip() ) first_num_match = re.match(r"^§(\d+)$", inside_content.strip()) # get all level 1 titles title_name = [ dictionary["first_title"] for dictionary in self.all_title if "first_title" in dictionary ] if first_level_match: first_title_text = first_level_match.group(2) first_title_num = first_level_match.group(1) first_title = first_title_num + first_title_text # the title does not contain "..." and is not in the title list # , add it to the title list if first_title not in title_name and ( int(first_title_num) == 1 or int(first_title_num) - int(self.all_title[-1]["id"]) == 1 ): current_entry = { "id": first_title_num, "first_title": first_title, "second_title": [], "table": [], } self.all_title.append(current_entry) elif second_level_match: second_title_name = second_level_match.group(0) second_title = second_level_match.group(1) first_title = second_title.split(".")[0] if (int(first_title) - 1 >= len(self.all_title)) or int( first_title ) - 1 < 0: continue else: titles = [ sub_item["title"] for sub_item in self.all_title[int(first_title) - 1][ "second_title" ] ] if second_title_name not in titles: self.all_title[int(first_title) - 1]["second_title"].append( {"title": second_title_name, "table": []} ) elif first_num_match: first_num = first_num_match.group(1) first_text = self.all_text[i + 1].get("inside") first_title = first_num_match.group(1) + first_text # if the title does not contain "..." and is not in the title list if ( "..." not in first_text and first_title not in title_name and ( int(first_num) == 1 or int(first_num) - int(self.all_title[-1]["id"]) == 1 ) ): current_entry = { "id": first_num, "first_title": first_title, "second_title": [], "table": [], } self.all_title.append(current_entry) def _load(self) -> List[Document]: """Load pdf document from loader.""" if self._loader: documents = self._loader.load() else: self._pdf_processor.pdf_to_json() file_title = self.file_path.rsplit("/", 1)[-1].replace(".pdf", "") self.all_text = list(self._pdf_processor.all_text.values()) self.process_text_data() temp_table = [] temp_title = None page_documents = [] merged_data = {} # type: ignore # noqa for i, data in enumerate(self.all_text): content_type = data.get("type") inside_content = data.get("inside") page = data.get("page") if content_type == "excel": temp_table.append(inside_content) if temp_title is None: for j in range(i - 1, -1, -1): if self.all_text[j]["type"] == "excel": break if self.all_text[j]["type"] == "text": content = self.all_text[j]["inside"] if re.match( r"^\d+\.\d+", content ) or content.startswith("§"): temp_title = content.strip() break else: temp_title = content.strip() break elif content_type == "text": if page in merged_data: # page merge merged_data[page]["inside_content"] += " " + inside_content else: merged_data[page] = { "inside_content": inside_content, "type": "text", } # merge excel table if temp_table: table_meta = { "title": temp_title or temp_table[0], "type": "excel", } self.all_title.append(table_meta) # markdown format markdown_tables = [] if temp_table: header = eval(temp_table[0]) markdown_tables.append(header) for entry in temp_table[1:]: row = eval(entry) markdown_tables.append(row) markdown_output = "| " + " | ".join(header) + " |\n" markdown_output += ( "| " + " | ".join(["---"] * len(header)) + " |\n" ) for row in markdown_tables[1:]: markdown_output += "| " + " | ".join(row) + " |\n" # merged content merged_data[page]["excel_content"] = temp_table merged_data[page]["markdown_output"] = markdown_output temp_title = None temp_table = [] # deal last excel if temp_table: table_meta = { "title": temp_title or temp_table[0], "table": temp_table, "type": "excel", } self.all_title.append(table_meta) # markdown format markdown_tables = [] if temp_table: header = eval(temp_table[0]) markdown_tables.append(header) for entry in temp_table[1:]: row = eval(entry) markdown_tables.append(row) markdown_output = "| " + " | ".join(header) + " |\n" markdown_output += "| " + " | ".join(["---"] * len(header)) + " |\n" for row in markdown_tables[1:]: markdown_output += "| " + " | ".join(row) + " |\n" # merged content merged_data[page]["excel_content"] = temp_table merged_data[page]["markdown_output"] = markdown_output for page, content in merged_data.items(): inside_content = content["inside_content"] if "markdown_output" in content: markdown_content = content["markdown_output"] content_metadata = { "page": page, "type": "excel", "title": file_title, "source": self.file_path, } page_documents.append( Document( content=inside_content + "\n" + markdown_content, metadata=content_metadata, ) ) else: content_metadata = { "page": page, "type": "text", "title": file_title, "source": self.file_path, } page_documents.append( Document(content=inside_content, metadata=content_metadata) ) return page_documents return [Document.langchain2doc(lc_document) for lc_document in documents] @classmethod def support_chunk_strategy(cls) -> List[ChunkStrategy]: """Return support chunk strategy.""" return [ ChunkStrategy.CHUNK_BY_SIZE, ChunkStrategy.CHUNK_BY_PAGE, ChunkStrategy.CHUNK_BY_SEPARATOR, ] @classmethod def default_chunk_strategy(cls) -> ChunkStrategy: """Return default chunk strategy.""" return ChunkStrategy.CHUNK_BY_SIZE @classmethod def type(cls) -> KnowledgeType: """Return knowledge type.""" return KnowledgeType.DOCUMENT @classmethod def document_type(cls) -> DocumentType: """Document type of PDF.""" return DocumentType.PDF class PDFProcessor: """PDFProcessor class.""" def __init__(self, filepath): """Initialize PDFProcessor class.""" self.filepath = filepath try: import pdfplumber # type: ignore except ImportError: raise ImportError("Please install pdfplumber first.") self.pdf = pdfplumber.open(filepath) self.all_text = defaultdict(dict) self.allrow = 0 self.last_num = 0 def check_lines(self, page, top, buttom): """Check lines.""" lines = page.extract_words()[::] text = "" last_top = 0 last_check = 0 for line in range(len(lines)): each_line = lines[line] check_re = ( "(?:。|；|单位：人民币元|金额单位：人民币元|单位：万元|币种：人民币|\d|" "报告(?:全文)?(?:（修订版）|（修订稿）|（更正后）)?)$" ) if top == "" and buttom == "": if abs(last_top - each_line["top"]) <= 2 or ( last_check > 0 and (page.height * 0.9 - each_line["top"]) > 0 and not re.search(check_re, text) ): text = text + each_line["text"] else: text = text + "\n" + each_line["text"] elif top == "": if each_line["top"] > buttom: if abs(last_top - each_line["top"]) <= 2 or ( last_check > 0 and (page.height * 0.85 - each_line["top"]) > 0 and not re.search(check_re, text) ): text = text + each_line["text"] else: text = text + "\n" + each_line["text"] else: if each_line["top"] < top and each_line["top"] > buttom: if abs(last_top - each_line["top"]) <= 2 or ( last_check > 0 and (page.height * 0.85 - each_line["top"]) > 0 and not re.search(check_re, text) ): text = text + each_line["text"] else: text = text + "\n" + each_line["text"] last_top = each_line["top"] last_check = each_line["x1"] - page.width * 0.85 return text def drop_empty_cols(self, data): """Delete empty column.""" transposed_data = list(map(list, zip(*data))) filtered_data = [ col for col in transposed_data if not all(cell == "" for cell in col) ] result = list(map(list, zip(*filtered_data))) return result def extract_text_and_tables(self, page): """Extract text and tables.""" buttom = 0 tables = page.find_tables() if len(tables) >= 1: count = len(tables) for table in tables: # process text before table if table.bbox[3] < buttom: pass else: count -= 1 # process text before table top = table.bbox[1] text = self.check_lines(page, top, buttom) text_list = text.split("\n") for _t in range(len(text_list)): self.all_text[self.allrow] = { "page": page.page_number, "allrow": self.allrow, "type": "text", "inside": text_list[_t], } self.allrow += 1 # process table buttom = table.bbox[3] new_table = table.extract() r_count = 0 for r in range(len(new_table)): row = new_table[r] if row[0] is None: r_count += 1 for c in range(len(row)): if row[c] is not None and row[c] not in ["", " "]: if new_table[r - r_count][c] is None: new_table[r - r_count][c] = row[c] else: new_table[r - r_count][c] += row[c] new_table[r][c] = None else: r_count = 0 end_table = [] for row in new_table: if row[0] is not None: cell_list = [] cell_check = False for cell in row: if cell is not None: cell = cell.replace("\n", "") else: cell = "" if cell != "": cell_check = True cell_list.append(cell) if cell_check: end_table.append(cell_list) end_table = self.drop_empty_cols(end_table) # process when column name is empty if len(end_table) > 0: for i in range(len(end_table[0])): if end_table[0][i] == "": if 0 < i < len(end_table[0]) - 1: # left column name left_column = end_table[0][i - 1] # right column name right_column = end_table[0][i + 1] # current name = left name + right name end_table[0][i] = left_column + right_column else: # if current column is empty and is the first # column, assign the right column name. # if current column is empty and is the # last column, assign the left column name. end_table[0][i] = ( end_table[0][i - 1] if i == len(end_table[0]) - 1 else end_table[0][i + 1] ) # if the first row is empty, assign the value of the previous row for i in range(1, len(end_table)): for j in range(len(end_table[i])): if end_table[i][j] == "": end_table[i][j] = end_table[i][j - 1] for row in end_table: self.all_text[self.allrow] = { "page": page.page_number, "allrow": self.allrow, "type": "excel", "inside": str(row), } self.allrow += 1 if count == 0: text = self.check_lines(page, "", buttom) text_list = text.split("\n") for _t in range(len(text_list)): self.all_text[self.allrow] = { "page": page.page_number, "allrow": self.allrow, "type": "text", "inside": text_list[_t], } self.allrow += 1 else: text = self.check_lines(page, "", "") text_list = text.split("\n") for _t in range(len(text_list)): self.all_text[self.allrow] = { "page": page.page_number, "allrow": self.allrow, "type": "text", "inside": text_list[_t], } self.allrow += 1 first_re = "[^计](?:报告(?:全文)?(?:（修订版）|（修订稿）|（更正后）)?)$" end_re = "^(?:\d|\\|\/|第|共|页|-|_| ){1,}" if self.last_num == 0: try: first_text = str(self.all_text[1]["inside"]) end_text = str(self.all_text[len(self.all_text) - 1]["inside"]) if re.search(first_re, first_text) and "[" not in end_text: self.all_text[1]["type"] = "页眉" if re.search(end_re, end_text) and "[" not in end_text: self.all_text[len(self.all_text) - 1]["type"] = "页脚" except Exception: print(page.page_number) else: try: first_text = str(self.all_text[self.last_num + 2]["inside"]) end_text = str(self.all_text[len(self.all_text) - 1]["inside"]) if re.search(first_re, first_text) and "[" not in end_text: self.all_text[self.last_num + 2]["type"] = "页眉" if re.search(end_re, end_text) and "[" not in end_text: self.all_text[len(self.all_text) - 1]["type"] = "页脚" except Exception: print(page.page_number) self.last_num = len(self.all_text) - 1 def pdf_to_json(self): """Process pdf.""" for i in range(len(self.pdf.pages)): self.extract_text_and_tables(self.pdf.pages[i]) logger.info(f"{self.filepath} page {i} extract text success") def save_all_text(self, path): """Save all text.""" directory = os.path.dirname(path) if not os.path.exists(directory): os.makedirs(directory) for key in self.all_text.keys(): with open(path, "a+", encoding="utf-8") as file: file.write(json.dumps(self.all_text[key], ensure_ascii=False) + "\n")