diff --git a/pageindex/utils.py b/pageindex/utils.py index 6563d26..178952f 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -391,9 +391,9 @@ SUPPORTED_PDF_PARSERS = ("PyPDF2", "pypdfium2", "PyMuPDF") DEFAULT_PDF_PARSER = os.getenv("PAGEINDEX_PDF_PARSER") or SUPPORTED_PDF_PARSERS[0] -def read_pdf_pages(doc): - """Return a list of per-page text strings using the currently configured parser.""" - parser = DEFAULT_PDF_PARSER +def read_pdf_pages(doc, pdf_parser=None): + """Return a list of per-page text strings using the selected parser.""" + parser = pdf_parser or DEFAULT_PDF_PARSER if parser == "PyPDF2": reader = PyPDF2.PdfReader(doc) @@ -435,15 +435,14 @@ def read_pdf_pages(doc): d.close() raise ValueError( - f"Unsupported DEFAULT_PDF_PARSER={parser!r}. Choose from {SUPPORTED_PDF_PARSERS}." + f"Unsupported pdf_parser={parser!r}. Choose from {SUPPORTED_PDF_PARSERS}." ) -def get_page_tokens(pdf_path, model=None): - pages = read_pdf_pages(pdf_path) +def get_page_tokens(pdf_path, model=None, pdf_parser=None): + pages = read_pdf_pages(pdf_path, pdf_parser=pdf_parser) return [(text, litellm.token_counter(model=model, text=text)) for text in pages] - def get_text_of_pdf_pages(pdf_pages, start_page, end_page): text = "" diff --git a/requirements.txt b/requirements.txt index e6ab063..70907ca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ litellm==1.83.7 # openai-agents # optional: required for examples/agentic_vectorless_rag_demo.py pymupdf==1.26.4 -# pypdfium2 # optional: enables pdf_parser="pypdfium2" (cleaner text, faster, Apache 2.0) +# pypdfium2 # optional: alternative PDF parser PyPDF2==3.0.1 python-dotenv==1.2.2 pyyaml==6.0.2 diff --git a/run_pageindex.py b/run_pageindex.py index 3548747..bd5efa4 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -29,7 +29,8 @@ if __name__ == "__main__": 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, - help='PDF text extractor: PyPDF2 (default), pypdfium2 (requires `pip install pypdfium2`), or PyMuPDF') + choices=pageindex_utils.SUPPORTED_PDF_PARSERS, + help='PDF text extractor to use') # Markdown specific arguments parser.add_argument('--if-thinning', type=str, default='no', @@ -136,4 +137,4 @@ if __name__ == "__main__": with open(output_file, 'w', encoding='utf-8') as f: json.dump(toc_with_page_number, f, indent=2, ensure_ascii=False) - print(f'Tree structure saved to: {output_file}') \ No newline at end of file + print(f'Tree structure saved to: {output_file}')