diff --git a/pageindex/config.yaml b/pageindex/config.yaml index 73a512c7a..6d65edbed 100644 --- a/pageindex/config.yaml +++ b/pageindex/config.yaml @@ -10,4 +10,8 @@ max_token_num_each_node: 20000 if_add_node_id: "yes" if_add_node_summary: "yes" if_add_doc_description: "no" -if_add_node_text: "no" \ No newline at end of file +if_add_node_text: "no" +# PDF text extractor. One of: PyPDF2 (default), PyMuPDF, pdf_inspector. +# pdf_inspector requires `pip install pdf-inspector` and preserves headings, +# lists and tables from text-based PDFs. +pdf_parser: "PyPDF2" \ No newline at end of file diff --git a/pageindex/page_index.py b/pageindex/page_index.py index c0b3ea935..12b9fd6a8 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -1240,7 +1240,7 @@ def page_index_main(doc, opt=None): raise ValueError("Unsupported input type. Expected a PDF file path or BytesIO object.") print('Parsing PDF...') - page_list = get_page_tokens(doc, model=opt.model) + page_list = get_page_tokens(doc, model=opt.model, pdf_parser=getattr(opt, 'pdf_parser', 'PyPDF2')) logger.info({'total_page_number': len(page_list)}) logger.info({'total_token': sum([page[1] for page in page_list])}) diff --git a/pageindex/utils.py b/pageindex/utils.py index 92fc46d85..1492ec4dc 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -20,6 +20,57 @@ # litellm is imported inside the functions that use it; eager import is slow # and fetches a remote model-cost map. +# pdf-inspector is an optional backend imported lazily so its absence does not +# break users on the default PyPDF2 path. +SUPPORTED_PDF_PARSERS = ("PyPDF2", "PyMuPDF", "pdf_inspector") + + +def _load_pdf_inspector(): + try: + import pdf_inspector # noqa: F401 + return pdf_inspector + except ImportError as e: + raise ImportError( + "pdf_parser='pdf_inspector' requires the 'pdf-inspector' package. " + "Install it with `pip install pdf-inspector`." + ) from e + + +def _pdf_inspector_pages(pdf_path, pages=None): + """Return the raw list of PageMarkdown items from pdf-inspector. + + Accepts a filesystem path or a BytesIO buffer. `pages` is an optional list + of 0-indexed page numbers matching pdf-inspector's convention. + """ + pi = _load_pdf_inspector() + if isinstance(pdf_path, BytesIO): + buf = pdf_path.getvalue() + return pi.extract_pages_markdown_bytes(buf, pages=pages).pages + return pi.extract_pages_markdown(pdf_path, pages=pages).pages + + +def classify_pdf(pdf_path): + """Classify a PDF as text_based / scanned / image_based / mixed. + + Returns a dict with `pdf_type`, `page_count`, `pages_needing_ocr` + (1-indexed), or None if pdf-inspector is not installed. Useful as a + preflight before running the full PageIndex pipeline. + """ + try: + pi = _load_pdf_inspector() + except ImportError: + return None + if isinstance(pdf_path, BytesIO): + result = pi.classify_pdf_bytes(pdf_path.getvalue()) + else: + result = pi.classify_pdf(pdf_path) + return { + "pdf_type": result.pdf_type, + "page_count": result.page_count, + # pdf-inspector reports 0-indexed pages; PageIndex uses 1-indexed. + "pages_needing_ocr": [p + 1 for p in getattr(result, "pages_needing_ocr", []) or []], + } + # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"): os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY") @@ -276,22 +327,62 @@ def get_last_node(structure): return structure[-1] -def extract_text_from_pdf(pdf_path): +def extract_text_from_pdf(pdf_path, pdf_parser="PyPDF2"): + if pdf_parser == "pdf_inspector": + pages = _pdf_inspector_pages(pdf_path) + return "".join((p.markdown or "") for p in pages) + if pdf_parser == "PyMuPDF": + if isinstance(pdf_path, BytesIO): + doc = pymupdf.open(stream=pdf_path, filetype="pdf") + else: + doc = pymupdf.open(pdf_path) + return "".join(page.get_text() for page in doc) + if pdf_parser != "PyPDF2": + raise ValueError(f"Unsupported PDF parser: {pdf_parser}") pdf_reader = PyPDF2.PdfReader(pdf_path) - ###return text not list + ###return text not list text="" for page_num in range(len(pdf_reader.pages)): page = pdf_reader.pages[page_num] text+=page.extract_text() return text -def get_pdf_title(pdf_path): +def get_pdf_title(pdf_path, pdf_parser="PyPDF2"): + # Title lives in the PDF metadata dictionary regardless of the text-parser + # backend; PyPDF2 handles that cheaply so we keep it as the single path. pdf_reader = PyPDF2.PdfReader(pdf_path) meta = pdf_reader.metadata title = meta.title if meta and meta.title else 'Untitled' return title -def get_text_of_pages(pdf_path, start_page, end_page, tag=True): +def get_text_of_pages(pdf_path, start_page, end_page, tag=True, pdf_parser="PyPDF2"): + if pdf_parser == "pdf_inspector": + # pdf-inspector uses 0-indexed page numbers. + wanted = list(range(start_page - 1, end_page)) + pages = _pdf_inspector_pages(pdf_path, pages=wanted) + text = "" + for page_num, page in zip(wanted, pages): + page_text = page.markdown or "" + if tag: + text += f"\n{page_text}\n\n" + else: + text += page_text + return text + if pdf_parser == "PyMuPDF": + if isinstance(pdf_path, BytesIO): + doc = pymupdf.open(stream=pdf_path, filetype="pdf") + else: + doc = pymupdf.open(pdf_path) + text = "" + for page_num in range(start_page - 1, end_page): + page_text = doc[page_num].get_text() + if tag: + text += f"\n{page_text}\n\n" + else: + text += page_text + return text + if pdf_parser != "PyPDF2": + raise ValueError(f"Unsupported PDF parser: {pdf_parser}") pdf_reader = PyPDF2.PdfReader(pdf_path) text = "" for page_num in range(start_page-1, end_page): @@ -465,6 +556,17 @@ def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"): token_length = litellm.token_counter(model=model, text=page_text) page_list.append((page_text, token_length)) return page_list + elif pdf_parser == "pdf_inspector": + # pdf-inspector emits per-page Markdown with heading tiers, list + # markers and GFM tables preserved. That structure carries into the + # downstream LLM prompt and node summaries. + pages = _pdf_inspector_pages(pdf_path) + page_list = [] + for page in pages: + page_text = page.markdown or "" + token_length = litellm.token_counter(model=model, text=page_text) + page_list.append((page_text, token_length)) + return page_list else: raise ValueError(f"Unsupported PDF parser: {pdf_parser}") @@ -482,7 +584,16 @@ def get_text_of_pdf_pages_with_labels(pdf_pages, start_page, end_page): text += f"\n{pdf_pages[page_num][0]}\n\n" return text -def get_number_of_pages(pdf_path): +def get_number_of_pages(pdf_path, pdf_parser="PyPDF2"): + if pdf_parser == "pdf_inspector": + try: + pi = _load_pdf_inspector() + except ImportError: + pass + else: + if isinstance(pdf_path, BytesIO): + return pi.classify_pdf_bytes(pdf_path.getvalue()).page_count + return pi.classify_pdf(pdf_path).page_count pdf_reader = PyPDF2.PdfReader(pdf_path) num = len(pdf_reader.pages) return num diff --git a/run_pageindex.py b/run_pageindex.py index 452f08174..3acc608eb 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -39,7 +39,17 @@ help='Whether to add doc description to the doc') parser.add_argument('--if-add-node-text', type=str, default=None, help='Whether to add text to the node') - + + parser.add_argument('--pdf-parser', type=str, default=None, + choices=['PyPDF2', 'PyMuPDF', 'pdf_inspector'], + help='PDF text extractor backend (PDF only, non-flash). ' + '`pdf_inspector` preserves headings, lists and tables ' + 'and requires `pip install pdf-inspector`.') + parser.add_argument('--check-ocr', action='store_true', + help='Classify the PDF with pdf-inspector before parsing ' + 'and refuse image-based files that would need OCR ' + '(requires `pip install pdf-inspector`).') + # Markdown specific arguments parser.add_argument('--if-thinning', type=str, default='no', help='Whether to apply tree thinning for markdown (markdown only)') @@ -67,8 +77,27 @@ raise ValueError("PDF file must have .pdf extension") if not os.path.isfile(args.pdf_path): raise ValueError(f"PDF file not found: {args.pdf_path}") - + + if args.check_ocr: + from pageindex.utils import classify_pdf + info = classify_pdf(args.pdf_path) + if info is None: + raise SystemExit( + "--check-ocr requires the `pdf-inspector` package. " + "Install it with `pip install pdf-inspector`." + ) + print(f"PDF type: {info['pdf_type']} ({info['page_count']} pages)") + if info['pdf_type'] in ('scanned', 'image_based'): + raise SystemExit( + f"Refusing to parse: this PDF is {info['pdf_type']} and " + "needs OCR. PageIndex does not run OCR locally." + ) + if info['pages_needing_ocr']: + print(f"Warning: pages needing OCR: {info['pages_needing_ocr']}") + if args.flash: + if args.pdf_parser is not None: + raise ValueError("--pdf-parser is not supported with --flash") from pageindex.flash import page_index_flash if args.optimize == 'full': from pageindex.tree_optimize import default_model @@ -102,6 +131,7 @@ 'if_add_node_summary': args.if_add_node_summary, 'if_add_doc_description': args.if_add_doc_description, 'if_add_node_text': args.if_add_node_text, + 'pdf_parser': args.pdf_parser, } opt = ConfigLoader().load({k: v for k, v in user_opt.items() if v is not None}) toc_with_page_number = page_index_main(args.pdf_path, opt) diff --git a/tests/test_pdf_parser_backend.py b/tests/test_pdf_parser_backend.py new file mode 100644 index 000000000..8a787f107 --- /dev/null +++ b/tests/test_pdf_parser_backend.py @@ -0,0 +1,130 @@ +import os +import sys +import unittest +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from pageindex import utils +from pageindex.utils import ( + SUPPORTED_PDF_PARSERS, + classify_pdf, + extract_text_from_pdf, + get_number_of_pages, + get_page_tokens, + get_text_of_pages, +) + + +FIXTURE_PDF = os.path.join( + os.path.dirname(__file__), + "..", + "examples", + "documents", + "q1-fy25-earnings.pdf", +) + + +def _has_pdf_inspector(): + try: + import pdf_inspector # noqa: F401 + return True + except ImportError: + return False + + +REQUIRES_PDF_INSPECTOR = unittest.skipUnless( + _has_pdf_inspector(), "pdf-inspector not installed" +) + + +class BackendDispatchTest(unittest.TestCase): + def test_supported_parsers_advertised(self): + self.assertIn("PyPDF2", SUPPORTED_PDF_PARSERS) + self.assertIn("PyMuPDF", SUPPORTED_PDF_PARSERS) + self.assertIn("pdf_inspector", SUPPORTED_PDF_PARSERS) + + def test_get_page_tokens_rejects_unknown_parser(self): + with self.assertRaises(ValueError): + get_page_tokens("dummy.pdf", model="gpt-4o", pdf_parser="not_a_parser") + + def test_extract_text_from_pdf_rejects_unknown_parser(self): + with self.assertRaises(ValueError): + extract_text_from_pdf("dummy.pdf", pdf_parser="not_a_parser") + + def test_get_text_of_pages_rejects_unknown_parser(self): + with self.assertRaises(ValueError): + get_text_of_pages("dummy.pdf", 1, 1, pdf_parser="not_a_parser") + + def test_load_pdf_inspector_raises_friendly_error_when_missing(self): + # Simulate the package being absent even when it is installed on the + # dev machine, so this test runs uniformly in CI. + with patch.dict(sys.modules, {"pdf_inspector": None}): + with self.assertRaises(ImportError) as ctx: + utils._load_pdf_inspector() + self.assertIn("pdf-inspector", str(ctx.exception)) + + def test_classify_pdf_returns_none_when_pdf_inspector_missing(self): + with patch.dict(sys.modules, {"pdf_inspector": None}): + self.assertIsNone(classify_pdf("dummy.pdf")) + + def test_get_number_of_pages_falls_back_when_pdf_inspector_missing(self): + # Backend requested but unavailable: fall back to PyPDF2 rather than + # raise, so opt.pdf_parser stays a soft preference. + fake_reader = MagicMock() + fake_reader.pages = [object(), object(), object()] + with patch.dict(sys.modules, {"pdf_inspector": None}), \ + patch("pageindex.utils.PyPDF2.PdfReader", return_value=fake_reader): + self.assertEqual( + get_number_of_pages("dummy.pdf", pdf_parser="pdf_inspector"), + 3, + ) + + +@REQUIRES_PDF_INSPECTOR +class PdfInspectorBackendTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + if not os.path.isfile(FIXTURE_PDF): + raise unittest.SkipTest(f"fixture PDF missing: {FIXTURE_PDF}") + + def test_classify_pdf_returns_expected_fields(self): + info = classify_pdf(FIXTURE_PDF) + self.assertIsNotNone(info) + self.assertIn("pdf_type", info) + self.assertIn("page_count", info) + self.assertIn("pages_needing_ocr", info) + self.assertGreater(info["page_count"], 0) + self.assertIn(info["pdf_type"], {"text_based", "mixed", "scanned", "image_based"}) + + def test_get_number_of_pages_matches_pypdf2(self): + expected = get_number_of_pages(FIXTURE_PDF, pdf_parser="PyPDF2") + actual = get_number_of_pages(FIXTURE_PDF, pdf_parser="pdf_inspector") + self.assertEqual(actual, expected) + + def test_get_page_tokens_returns_all_pages(self): + pages = get_page_tokens(FIXTURE_PDF, model="gpt-4o", pdf_parser="pdf_inspector") + self.assertGreater(len(pages), 0) + expected = get_number_of_pages(FIXTURE_PDF, pdf_parser="PyPDF2") + self.assertEqual(len(pages), expected) + # Every entry must be (text, token_count). + for text, tokens in pages: + self.assertIsInstance(text, str) + self.assertIsInstance(tokens, int) + + def test_get_text_of_pages_tags_selected_range(self): + text = get_text_of_pages( + FIXTURE_PDF, 3, 4, tag=True, pdf_parser="pdf_inspector" + ) + self.assertIn("", text) + self.assertIn("", text) + self.assertIn("", text) + self.assertIn("", text) + + def test_extract_text_from_pdf_returns_nonempty(self): + text = extract_text_from_pdf(FIXTURE_PDF, pdf_parser="pdf_inspector") + self.assertGreater(len(text), 100) + + +if __name__ == "__main__": + unittest.main()