# # Copyright 2025 The InfiniFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import logging from io import BytesIO import re from deepdoc.parser.utils import get_text from rag.app import naive from rag.nlp import rag_tokenizer, tokenize from deepdoc.parser import PdfParser, ExcelParser, HtmlParser from deepdoc.parser.figure_parser import vision_figure_parser_docx_wrapper_naive from rag.app.naive import by_plaintext, PARSERS from common.constants import MAXIMUM_PAGE_NUMBER, MAXIMUM_TASK_PAGE_NUMBER from common.parser_config_utils import normalize_layout_recognizer class Pdf(PdfParser): def __call__(self, filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, zoomin=3, callback=None): from timeit import default_timer as timer start = timer() callback(msg="OCR started") self.__images__(filename if not binary else binary, zoomin, from_page, to_page, callback) callback(msg="OCR finished ({:.2f}s)".format(timer() - start)) start = timer() self._layouts_rec(zoomin, drop=False) callback(0.63, "Layout analysis ({:.2f}s)".format(timer() - start)) logging.debug("layouts cost: {}s".format(timer() - start)) start = timer() self._table_transformer_job(zoomin) callback(0.65, "Table analysis ({:.2f}s)".format(timer() - start)) start = timer() self._text_merge() callback(0.67, "Text merged ({:.2f}s)".format(timer() - start)) tbls = self._extract_table_figure(True, zoomin, True, True) self._concat_downward() sections = [(b["text"], self.get_position(b, zoomin)) for i, b in enumerate(self.boxes)] return [(txt, "") for txt, _ in sorted(sections, key=lambda x: (x[-1][0][0], x[-1][0][3], x[-1][0][1]))], tbls def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang="Chinese", callback=None, **kwargs): """ Supported file formats are docx, pdf, excel, txt. One file forms a chunk which maintains original text order. """ parser_config = kwargs.get("parser_config", {"chunk_token_num": 512, "delimiter": "\n!?。；！？", "layout_recognize": "DeepDOC"}) eng = lang.lower() == "english" # is_english(cks) if re.search(r"\.docx$", filename, re.IGNORECASE): callback(0.1, "Start to parse.") sections = naive.Docx()(filename, binary) cks = [] image_idxs = [] for text, image, table in sections: if table is not None: text = (text or "") + str(table) ck_type = "table" else: ck_type = "image" if image is not None else "text" if ck_type == "image": image_idxs.append(len(cks)) cks.append({"text": text, "image": image, "ck_type": ck_type}) vision_figure_parser_docx_wrapper_naive(cks, image_idxs, callback, **kwargs) sections = [ck["text"] for ck in cks if ck.get("text")] callback(0.8, "Finish parsing.") elif re.search(r"\.pdf$", filename, re.IGNORECASE): layout_recognizer, parser_model_name = normalize_layout_recognizer(parser_config.get("layout_recognize", "DeepDOC")) if isinstance(layout_recognizer, bool): layout_recognizer = "DeepDOC" if layout_recognizer else "Plain Text" name = layout_recognizer.strip().lower() parser = PARSERS.get(name, by_plaintext) callback(0.1, "Start to parse.") sections, tbls, pdf_parser = parser( filename=filename, binary=binary, from_page=from_page, to_page=to_page, lang=lang, callback=callback, pdf_cls=Pdf, layout_recognizer=layout_recognizer, mineru_llm_name=parser_model_name, paddleocr_llm_name=parser_model_name, **kwargs, ) if not sections and not tbls: return [] if name in ["tcadp", "docling", "mineru", "paddleocr"]: parser_config["chunk_token_num"] = 0 callback(0.8, "Finish parsing.") for (img, rows), poss in tbls: if not rows: continue sections.append((rows if isinstance(rows, str) else rows[0], [(p[0] + 1 - from_page, p[1], p[2], p[3], p[4]) for p in poss])) sections = [s for s, _ in sections if s] elif re.search(r"\.xlsx?$", filename, re.IGNORECASE): callback(0.1, "Start to parse.") excel_parser = ExcelParser() sections = excel_parser.html(binary, MAXIMUM_TASK_PAGE_NUMBER) elif re.search(r"\.(txt|md|markdown|mdx)$", filename, re.IGNORECASE): callback(0.1, "Start to parse.") txt = get_text(filename, binary) sections = txt.split("\n") sections = [s for s in sections if s] callback(0.8, "Finish parsing.") elif re.search(r"\.(htm|html)$", filename, re.IGNORECASE): callback(0.1, "Start to parse.") sections = HtmlParser()(filename, binary) sections = [s for s in sections if s] callback(0.8, "Finish parsing.") elif re.search(r"\.doc$", filename, re.IGNORECASE): callback(0.1, "Start to parse.") try: from tika import parser as tika_parser except Exception as e: callback(0.8, f"tika not available: {e}. Unsupported .doc parsing.") logging.warning(f"tika not available: {e}. Unsupported .doc parsing for {filename}.") return [] binary = BytesIO(binary) doc_parsed = tika_parser.from_buffer(binary) if doc_parsed.get("content", None) is not None: sections = doc_parsed["content"].split("\n") sections = [s for s in sections if s] callback(0.8, "Finish parsing.") else: raise NotImplementedError("file type not supported yet(doc, docx, pdf, txt supported)") doc = {"docnm_kwd": filename, "title_tks": rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", filename))} doc["title_sm_tks"] = rag_tokenizer.fine_grained_tokenize(doc["title_tks"]) tokenize(doc, "\n".join(sections), eng) return [doc] if __name__ == "__main__": import sys def dummy(prog=None, msg=""): pass chunk(sys.argv[1], from_page=0, to_page=10, callback=dummy)