Refactoring PDF loaders: 02 PyMuPDF

libs/community/extended_testing_deps.txt

@@ -66,6 +66,7 @@ praw>=7.7.1,<8
 premai>=0.3.25,<0.4
 psychicapi>=0.8.0,<0.9
 pydantic>=2.7.4,<3
+pytesseract>=0.3.13
 py-trello>=0.19.0,<0.20
 pyjwt>=2.8.0,<3
 pymupdf>=1.22.3,<2

libs/community/langchain_community/document_loaders/parsers/__init__.py

@@ -17,6 +17,11 @@
     from langchain_community.document_loaders.parsers.html import (
         BS4HTMLParser,
     )
+    from langchain_community.document_loaders.parsers.images import (
+        MultimodalBlobParser,
+        RapidOCRBlobParser,
+        TesseractBlobParser,
+    )
     from langchain_community.document_loaders.parsers.language import (
         LanguageParser,
     )
@@ -38,12 +43,15 @@
     "DocAIParser": "langchain_community.document_loaders.parsers.docai",
     "GrobidParser": "langchain_community.document_loaders.parsers.grobid",
     "LanguageParser": "langchain_community.document_loaders.parsers.language",
+    "MultimodalBlobParser": "langchain_community.document_loaders.parsers.images",
     "OpenAIWhisperParser": "langchain_community.document_loaders.parsers.audio",
     "PDFMinerParser": "langchain_community.document_loaders.parsers.pdf",
     "PDFPlumberParser": "langchain_community.document_loaders.parsers.pdf",
     "PyMuPDFParser": "langchain_community.document_loaders.parsers.pdf",
     "PyPDFParser": "langchain_community.document_loaders.parsers.pdf",
     "PyPDFium2Parser": "langchain_community.document_loaders.parsers.pdf",
+    "RapidOCRBlobParser": "langchain_community.document_loaders.parsers.images",
+    "TesseractBlobParser": "langchain_community.document_loaders.parsers.images",
     "VsdxParser": "langchain_community.document_loaders.parsers.vsdx",
 }
 
@@ -61,11 +69,14 @@ def __getattr__(name: str) -> Any:
     "DocAIParser",
     "GrobidParser",
     "LanguageParser",
+    "MultimodalBlobParser",
     "OpenAIWhisperParser",
     "PDFMinerParser",
     "PDFPlumberParser",
     "PyMuPDFParser",
     "PyPDFParser",
     "PyPDFium2Parser",
+    "RapidOCRBlobParser",
+    "TesseractBlobParser",
     "VsdxParser",
 ]

libs/community/langchain_community/document_loaders/parsers/images.py

@@ -0,0 +1,149 @@
+import base64
+import html
+import io
+import logging
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Iterator, Literal
+
+import numpy as np
+from langchain_core.documents import Document
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import HumanMessage
+
+if TYPE_CHECKING:
+    from PIL.Image import Image
+
+from langchain_community.document_loaders.base import BaseBlobParser
+from langchain_community.document_loaders.blob_loaders import Blob
+
+logger = logging.getLogger(__name__)
+
+
+class ImageBlobParser(BaseBlobParser):
+    def __init__(
+        self,
+        *,
+        format: Literal["text", "markdown", "html"] = "text",
+    ):
+        self.format = format
+
+    @abstractmethod
+    def _analyze_image(self, img: "Image") -> str:
+        pass
+
+    def lazy_parse(self, blob: Blob) -> Iterator[Document]:
+        try:
+            from PIL import Image as Img
+
+            with blob.as_bytes_io() as buf:
+                img = Img.open(buf)
+                content = self._analyze_image(img)
+                if content:
+                    if self.format == "markdown":
+                        content = content.replace("]", r"\\]")
+                        content = f"![{content}](.)"
+                    elif self.format == "html":
+                        content = f'<img alt="{html.escape(content, quote=True)}" />'
+                logger.debug("Image text: %s", content.replace("\n", "\\n"))
+                yield Document(
+                    page_content=content,
+                    metadata={"source": blob.source},
+                )
+        except ImportError:
+            raise ImportError(
+                "`rapidocr-onnxruntime` package not found, please install it with "
+                "`pip install Pillow`"
+            )
+
+
+class RapidOCRBlobParser(ImageBlobParser):
+    def __init__(
+        self,
+        *,
+        format: Literal["text", "markdown", "html"] = "text",
+    ):
+        super().__init__(format=format)
+        self.ocr = None
+
+    def _analyze_image(self, img: "Image") -> str:
+        if not self.ocr:
+            try:
+                from rapidocr_onnxruntime import RapidOCR
+
+                self.ocr = RapidOCR()
+            except ImportError:
+                raise ImportError(
+                    "`rapidocr-onnxruntime` package not found, please install it with "
+                    "`pip install rapidocr-onnxruntime`"
+                )
+        ocr_result, _ = self.ocr(np.array(img))  # type: ignore
+        content = ""
+        if ocr_result:
+            content = ("\n".join([text[1] for text in ocr_result])).strip()
+        return content
+
+
+class TesseractBlobParser(ImageBlobParser):
+    def __init__(
+        self,
+        *,
+        format: Literal["text", "markdown", "html"] = "text",
+        langs: list[str] = ["eng"],
+    ):
+        super().__init__(format=format)
+        self.langs = langs
+
+    def _analyze_image(self, img: "Image") -> str:
+        try:
+            import pytesseract
+        except ImportError:
+            raise ImportError(
+                "`pytesseract` package not found, please install it with "
+                "`pip install pytesseract`"
+            )
+        return pytesseract.image_to_string(img, lang="+".join(self.langs)).strip()
+
+
+_prompt_images_to_description = (
+    "You are an assistant tasked with summarizing "
+    "images for retrieval. "
+    "These summaries will be embedded and used to retrieve the raw image. "
+    "Give a concise summary of the image that is well optimized for retrieval "
+    "and extract all the text from the image."
+)
+
+
+class MultimodalBlobParser(ImageBlobParser):
+    def __init__(
+        self,
+        *,
+        format: Literal["text", "markdown", "html"] = "text",
+        model: BaseChatModel,
+        prompt: str = _prompt_images_to_description,
+    ):
+        super().__init__(format=format)
+        self.model = model
+        self.prompt = prompt
+
+    def _analyze_image(self, img: "Image") -> str:
+        image_bytes = io.BytesIO()
+        img.save(image_bytes, format="PNG")
+        img_base64 = base64.b64encode(image_bytes.getvalue()).decode("utf-8")
+        msg = self.model.invoke(
+            [
+                HumanMessage(
+                    content=[
+                        {"type": "text", "text": self.prompt},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{img_base64}"
+                            },
+                        },
+                    ]
+                )
+            ]
+        )
+        result = msg.content
+        assert isinstance(result, str)
+        return result