From 23231d65a902a476543ad4976b9d7f11d21994f7 Mon Sep 17 00:00:00 2001 From: Tim Asp <707699+timothyasp@users.noreply.github.com> Date: Fri, 3 Mar 2023 20:59:28 -0800 Subject: [PATCH] Add PyMuPDF PDF loader (#1426) Different PDF libraries have different strengths and weaknesses. PyMuPDF does a good job at extracting the most amount of content from the doc, regardless of the source quality, extremely fast (especially compared to Unstructured). https://pymupdf.readthedocs.io/en/latest/index.html --- .../document_loaders/examples/pdf.ipynb | 88 +++++++++++++++++-- langchain/document_loaders/__init__.py | 7 +- langchain/document_loaders/pdf.py | 37 +++++++- .../document_loaders/test_pdf.py | 46 ++++++++++ 4 files changed, 169 insertions(+), 9 deletions(-) create mode 100644 tests/integration_tests/document_loaders/test_pdf.py diff --git a/docs/modules/document_loaders/examples/pdf.ipynb b/docs/modules/document_loaders/examples/pdf.ipynb index 0e448ed2..f40d4807 100644 --- a/docs/modules/document_loaders/examples/pdf.ipynb +++ b/docs/modules/document_loaders/examples/pdf.ipynb @@ -104,10 +104,10 @@ "Efficient Data AnnotationC u s t o m i z e d M o d e l T r a i n i n gModel Cust omizationDI A Model HubDI A Pipeline SharingCommunity PlatformLa y out Detection ModelsDocument Images \n", "T h e C o r e L a y o u t P a r s e r L i b r a r yOCR ModuleSt or age & VisualizationLa y out Data Structur e\n", "Fig. 1: The overall architecture of LayoutParser . For an input document image,\n", - "the core LayoutParser library provides a set of o\u000b", + "the core LayoutParser library provides a set of o\u000B", "-the-shelf tools for layout\n", "detection, OCR, visualization, and storage, backed by a carefully designed layout\n", - "data structure. LayoutParser also supports high level customization via e\u000ecient\n", + "data structure. LayoutParser also supports high level customization via e\u000Ecient\n", "layout annotation and model training functions. These improve model accuracy\n", "on the target samples. The community platform enables the easy sharing of DIA\n", "models and whole digitization pipelines to promote reusability and reproducibility.\n", @@ -128,10 +128,10 @@ "gure layouts) and\n", "HJDataset [31](historical Japanese document layouts). A spectrum of models\n", "trained on these datasets are currently available in the LayoutParser model zoo\n", - "to support di\u000b", + "to support di\u000B", "erent use cases.\n", "3 The Core LayoutParser Library\n", - "At the core of LayoutParser is an o\u000b", + "At the core of LayoutParser is an o\u000B", "-the-shelf toolkit that streamlines DL-\n", "based document image analysis. Five components support a simple interface\n", "with comprehensive functionalities: 1) The layout detection models enable using\n", @@ -266,13 +266,87 @@ "data = loader.load()" ] }, + { + "cell_type": "markdown", + "source": [ + "## Using PyMuPDF\n", + "\n", + "This is the fastest of the PDF parsing options, and contains detailed metadata about the PDF and its pages, as well as returns one document per page." + ], + "metadata": { + "collapsed": false + } + }, + { + "cell_type": "code", + "execution_count": 1, + "outputs": [], + "source": [ + "from langchain.document_loaders import PyMuPDFLoader" + ], + "metadata": { + "collapsed": false + } + }, + { + "cell_type": "code", + "execution_count": 2, + "outputs": [], + "source": [ + "loader = PyMuPDFLoader(\"example_data/layout-parser-paper.pdf\")" + ], + "metadata": { + "collapsed": false + } + }, + { + "cell_type": "code", + "execution_count": 3, + "outputs": [], + "source": [ + "data = loader.load()" + ], + "metadata": { + "collapsed": false + } + }, + { + "cell_type": "code", + "execution_count": 4, + "outputs": [ + { + "data": { + "text/plain": "Document(page_content='LayoutParser: A Unified Toolkit for Deep\\nLearning Based Document Image Analysis\\nZejiang Shen1 (�), Ruochen Zhang2, Melissa Dell3, Benjamin Charles Germain\\nLee4, Jacob Carlson3, and Weining Li5\\n1 Allen Institute for AI\\nshannons@allenai.org\\n2 Brown University\\nruochen zhang@brown.edu\\n3 Harvard University\\n{melissadell,jacob carlson}@fas.harvard.edu\\n4 University of Washington\\nbcgl@cs.washington.edu\\n5 University of Waterloo\\nw422li@uwaterloo.ca\\nAbstract. Recent advances in document image analysis (DIA) have been\\nprimarily driven by the application of neural networks. Ideally, research\\noutcomes could be easily deployed in production and extended for further\\ninvestigation. However, various factors like loosely organized codebases\\nand sophisticated model configurations complicate the easy reuse of im-\\nportant innovations by a wide audience. Though there have been on-going\\nefforts to improve reusability and simplify deep learning (DL) model\\ndevelopment in disciplines like natural language processing and computer\\nvision, none of them are optimized for challenges in the domain of DIA.\\nThis represents a major gap in the existing toolkit, as DIA is central to\\nacademic research across a wide range of disciplines in the social sciences\\nand humanities. This paper introduces LayoutParser, an open-source\\nlibrary for streamlining the usage of DL in DIA research and applica-\\ntions. The core LayoutParser library comes with a set of simple and\\nintuitive interfaces for applying and customizing DL models for layout de-\\ntection, character recognition, and many other document processing tasks.\\nTo promote extensibility, LayoutParser also incorporates a community\\nplatform for sharing both pre-trained models and full document digiti-\\nzation pipelines. We demonstrate that LayoutParser is helpful for both\\nlightweight and large-scale digitization pipelines in real-word use cases.\\nThe library is publicly available at https://layout-parser.github.io.\\nKeywords: Document Image Analysis · Deep Learning · Layout Analysis\\n· Character Recognition · Open Source library · Toolkit.\\n1\\nIntroduction\\nDeep Learning(DL)-based approaches are the state-of-the-art for a wide range of\\ndocument image analysis (DIA) tasks including document image classification [11,\\narXiv:2103.15348v2 [cs.CV] 21 Jun 2021\\n', lookup_str='', metadata={'file_path': 'example_data/layout-parser-paper.pdf', 'page_number': 1, 'total_pages': 16, 'format': 'PDF 1.5', 'title': '', 'author': '', 'subject': '', 'keywords': '', 'creator': 'LaTeX with hyperref', 'producer': 'pdfTeX-1.40.21', 'creationDate': 'D:20210622012710Z', 'modDate': 'D:20210622012710Z', 'trapped': '', 'encryption': None}, lookup_index=0)" + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "data[0]" + ], + "metadata": { + "collapsed": false + } + }, + { + "cell_type": "markdown", + "source": [ + "Additionally, you can pass along any of the options from the [PyMuPDF documentation](https://pymupdf.readthedocs.io/en/latest/app1.html#plain-text/) as keyword arguments in the `load` call, and it will be pass along to the `get_text()` call." + ], + "metadata": { + "collapsed": false + } + }, { "cell_type": "code", "execution_count": null, - "id": "7301c473", - "metadata": {}, "outputs": [], - "source": [] + "source": [], + "metadata": { + "collapsed": false + } } ], "metadata": { diff --git a/langchain/document_loaders/__init__.py b/langchain/document_loaders/__init__.py index 8b1e722e..33710afc 100644 --- a/langchain/document_loaders/__init__.py +++ b/langchain/document_loaders/__init__.py @@ -24,7 +24,11 @@ from langchain.document_loaders.notion import NotionDirectoryLoader from langchain.document_loaders.obsidian import ObsidianLoader from langchain.document_loaders.online_pdf import OnlinePDFLoader from langchain.document_loaders.paged_pdf import PagedPDFSplitter -from langchain.document_loaders.pdf import PDFMinerLoader, UnstructuredPDFLoader +from langchain.document_loaders.pdf import ( + PDFMinerLoader, + PyMuPDFLoader, + UnstructuredPDFLoader, +) from langchain.document_loaders.powerpoint import UnstructuredPowerPointLoader from langchain.document_loaders.readthedocs import ReadTheDocsLoader from langchain.document_loaders.roam import RoamLoader @@ -78,6 +82,7 @@ __all__ = [ "AirbyteJSONLoader", "OnlinePDFLoader", "PDFMinerLoader", + "PyMuPDFLoader", "TelegramChatLoader", "SRTLoader", "FacebookChatLoader", diff --git a/langchain/document_loaders/pdf.py b/langchain/document_loaders/pdf.py index ca7b7e58..915b9fc8 100644 --- a/langchain/document_loaders/pdf.py +++ b/langchain/document_loaders/pdf.py @@ -1,5 +1,5 @@ """Loader that loads PDF files.""" -from typing import List +from typing import Any, List, Optional from langchain.docstore.document import Document from langchain.document_loaders.base import BaseLoader @@ -27,6 +27,7 @@ class PDFMinerLoader(BaseLoader): "pdfminer package not found, please install it with " "`pip install pdfminer.six`" ) + self.file_path = file_path def load(self) -> List[Document]: @@ -36,3 +37,37 @@ class PDFMinerLoader(BaseLoader): text = extract_text(self.file_path) metadata = {"source": self.file_path} return [Document(page_content=text, metadata=metadata)] + + +class PyMuPDFLoader(BaseLoader): + """Loader that uses PyMuPDF to load PDF files.""" + + def __init__(self, file_path: str): + """Initialize with file path.""" + try: + import fitz # noqa:F401 + except ImportError: + raise ValueError( + "PyMuPDF package not found, please install it with " + "`pip install pymupdf`" + ) + + self.file_path = file_path + + def load(self, **kwargs: Optional[Any]) -> List[Document]: + """Load file.""" + import fitz + + doc = fitz.open(self.file_path) # open document + return [ + Document( + page_content=page.get_text(**kwargs).encode("utf-8"), + metadata={ + "file_path": self.file_path, + "page_number": page.number + 1, + "total_pages": len(doc), + } + | doc.metadata, + ) + for page in doc + ] diff --git a/tests/integration_tests/document_loaders/test_pdf.py b/tests/integration_tests/document_loaders/test_pdf.py new file mode 100644 index 00000000..8e59f310 --- /dev/null +++ b/tests/integration_tests/document_loaders/test_pdf.py @@ -0,0 +1,46 @@ +from pathlib import Path + +from langchain.document_loaders import ( + PDFMinerLoader, + PyMuPDFLoader, + UnstructuredPDFLoader, +) + + +def test_unstructured_pdf_loader() -> None: + """Test unstructured loader.""" + file_path = Path(__file__).parent.parent / "examples/hello.pdf" + loader = UnstructuredPDFLoader(str(file_path)) + docs = loader.load() + + assert len(docs) == 1 + + +def test_pdfminer_loader() -> None: + """Test PDFMiner loader.""" + file_path = Path(__file__).parent.parent / "examples/hello.pdf" + loader = PDFMinerLoader(str(file_path)) + docs = loader.load() + + assert len(docs) == 1 + + file_path = Path(__file__).parent.parent / "examples/layout-parser-paper.pdf" + loader = PDFMinerLoader(str(file_path)) + + docs = loader.load() + assert len(docs) == 1 + + +def test_pymupdf_loader() -> None: + """Test PyMuPDF loader.""" + file_path = Path(__file__).parent.parent / "examples/hello.pdf" + loader = PyMuPDFLoader(str(file_path)) + + docs = loader.load() + assert len(docs) == 1 + + file_path = Path(__file__).parent.parent / "examples/layout-parser-paper.pdf" + loader = PyMuPDFLoader(str(file_path)) + + docs = loader.load() + assert len(docs) == 16