Add tests for llm processors

README.md

@@ -149,7 +149,7 @@ text, _, images = text_from_rendered(rendered)
 
 ### Custom configuration
 
-You can also pass configuration using the `ConfigParser`:
+You can pass configuration using the `ConfigParser`:
 
 ```python
 from marker.converters.pdf import PdfConverter
@@ -171,6 +171,26 @@ converter = PdfConverter(
 rendered = converter("FILEPATH")
 ```
 
+### Extract blocks
+
+Each document consists of one or more pages.  Pages contain blocks, which can themselves contain other blocks.  It's possible to programatically manipulate these blocks.  
+
+Here's an example of extracting all forms from a document:
+
+```python
+from marker.converters.pdf import PdfConverter
+from marker.models import create_model_dict
+from marker.schema import BlockTypes
+
+converter = PdfConverter(
+    artifact_dict=create_model_dict(),
+)
+document = converter.build_document("FILEPATH")
+forms = document.contained_blocks((BlockTypes.Form,))
+```
+
+Look at the processors for more examples of extracting and manipulating blocks.
+
 # Output Formats
 
 ## Markdown

convert.py

@@ -1,5 +1,7 @@
 import os
 
+os.environ["GRPC_VERBOSITY"] = "ERROR"
+os.environ["GLOG_minloglevel"] = "2"
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # Transformers uses .isin for a simple op, which is not supported on MPS
 os.environ["IN_STREAMLIT"] = "true" # Avoid multiprocessing inside surya
 

convert_single.py

@@ -1,9 +1,10 @@
 import os
 
+os.environ["GRPC_VERBOSITY"] = "ERROR"
+os.environ["GLOG_minloglevel"] = "2"
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"  # Transformers uses .isin for a simple op, which is not supported on MPS
 
 import time
-
 import click
 
 from marker.config.parser import ConfigParser

marker/builders/llm_layout.py

@@ -51,7 +51,7 @@ class LLMLayoutBuilder(LayoutBuilder):
     """
 
     google_api_key: Optional[str] = settings.GOOGLE_API_KEY
-    confidence_threshold: float = 0.7
+    confidence_threshold: float = 0.75
     model_name: str = "gemini-1.5-flash"
     max_retries: int = 3
     max_concurrency: int = 3

marker/converters/pdf.py

@@ -109,7 +109,7 @@ class PdfConverter(BaseConverter):
 
         return cls(**resolved_kwargs)
 
-    def __call__(self, filepath: str):
+    def build_document(self, filepath: str):
         pdf_provider = PdfProvider(filepath, self.config)
         layout_builder = self.resolve_dependencies(self.layout_builder_class)
         ocr_builder = self.resolve_dependencies(OcrBuilder)
@@ -120,5 +120,9 @@ class PdfConverter(BaseConverter):
             processor = self.resolve_dependencies(processor_cls)
             processor(document)
 
+        return document
+
+    def __call__(self, filepath: str):
+        document = self.build_document(filepath)
         renderer = self.resolve_dependencies(self.renderer)
         return renderer(document)

marker/processors/llm/__init__.py

@@ -32,6 +32,9 @@ class BaseLLMProcessor(BaseProcessor):
         gemini_rewriting_prompt (str):
             The prompt to use for rewriting text.
             Default is a string containing the Gemini rewriting prompt.
+        use_llm (bool):
+            Whether to use the LLM model.
+            Default is False.
     """
 
     google_api_key: Optional[str] = settings.GOOGLE_API_KEY
@@ -57,7 +60,10 @@ class BaseLLMProcessor(BaseProcessor):
         if not self.use_llm or self.model is None:
             return
 
-        self.rewrite_blocks(document)
+        try:
+            self.rewrite_blocks(document)
+        except Exception as e:
+            print(f"Error rewriting blocks in {self.__class__.__name__}: {e}")
 
     def process_rewriting(self, document: Document, page: PageGroup, block: Block):
         raise NotImplementedError()

marker/processors/llm/llm_form.py

@@ -1,11 +1,8 @@
 import markdown2
 
 from marker.processors.llm import BaseLLMProcessor
-from marker.processors.llm.utils import GoogleModel
-from concurrent.futures import ThreadPoolExecutor, as_completed
 
 from google.ai.generativelanguage_v1beta.types import content
-from tqdm import tqdm
 from tabled.formats import markdown_format
 
 from marker.schema import BlockTypes

marker/schema/document.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import List
+from typing import List, Sequence
 
 from pydantic import BaseModel
 
@@ -100,3 +100,9 @@ class Document(BaseModel):
             children=child_content,
             html=self.assemble_html(child_content)
         )
+
+    def contained_blocks(self, block_types: Sequence[BlockTypes] = None) -> List[Block]:
+        blocks = []
+        for page in self.pages:
+            blocks += page.contained_blocks(self, block_types)
+        return blocks

poetry.lock

@@ -3449,6 +3449,23 @@ tomli = {version = ">=1", markers = "python_version < \"3.11\""}
 [package.extras]
 dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"]
 
+[[package]]
+name = "pytest-mock"
+version = "3.14.0"
+description = "Thin-wrapper around the mock package for easier use with pytest"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "pytest-mock-3.14.0.tar.gz", hash = "sha256:2719255a1efeceadbc056d6bf3df3d1c5015530fb40cf347c0f9afac88410bd0"},
+    {file = "pytest_mock-3.14.0-py3-none-any.whl", hash = "sha256:0b72c38033392a5f4621342fe11e9219ac11ec9d375f8e2a0c164539e0d70f6f"},
+]
+
+[package.dependencies]
+pytest = ">=6.2.5"
+
+[package.extras]
+dev = ["pre-commit", "pytest-asyncio", "tox"]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"
@@ -5407,4 +5424,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "20eee90138195d778e93da276c2d02e6547738e8eedf3c0a355eaecb128a58c0"
+content-hash = "2a4dfa94c63b5cf4b614fb4908abd2c80c363e9ed4ebf53b71af9bba90b783fd"

pyproject.toml

@@ -50,6 +50,7 @@ fastapi = "^0.115.4"
 uvicorn = "^0.32.0"
 python-multipart = "^0.0.16"
 pytest = "^8.3.3"
+pytest-mock = "^3.14.0"
 
 [tool.poetry.scripts]
 marker = "convert:main"

tests/processors/test_llm_processors.py

@@ -0,0 +1,95 @@
+from unittest.mock import MagicMock, Mock
+
+import pytest
+
+from marker.processors.llm.llm_form import LLMFormProcessor
+from marker.processors.llm.llm_table import LLMTableProcessor
+from marker.processors.llm.llm_text import LLMTextProcessor
+from marker.processors.table import TableProcessor
+from marker.schema import BlockTypes
+
+@pytest.mark.filename("form_1040.pdf")
+@pytest.mark.config({"page_range": [0]})
+def test_llm_form_processor_no_config(pdf_document):
+    processor = LLMFormProcessor()
+    processor(pdf_document)
+
+    forms = pdf_document.contained_blocks((BlockTypes.Form,))
+    assert forms[0].html is None
+
+
+@pytest.mark.filename("form_1040.pdf")
+@pytest.mark.config({"page_range": [0]})
+def test_llm_form_processor_no_cells(pdf_document):
+    processor = LLMFormProcessor({"use_llm": True})
+    processor(pdf_document)
+
+    forms = pdf_document.contained_blocks((BlockTypes.Form,))
+    assert forms[0].html is None
+
+
+@pytest.mark.filename("form_1040.pdf")
+@pytest.mark.config({"page_range": [0]})
+def test_llm_form_processor(pdf_document, detection_model, table_rec_model, recognition_model, mocker):
+    corrected_markdown = "*This is corrected markdown.*\n" * 100
+
+    corrected_html = "<em>This is corrected markdown.</em>\n" * 100
+    corrected_html = "<p>" + corrected_html.strip() + "</p>\n"
+
+    mock_cls = Mock()
+    mock_cls.return_value.generate_response.return_value = {"corrected_markdown": corrected_markdown}
+    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+
+    cell_processor = TableProcessor(detection_model, recognition_model, table_rec_model)
+    cell_processor(pdf_document)
+
+    processor = LLMFormProcessor({"use_llm": True})
+    processor(pdf_document)
+
+    forms = pdf_document.contained_blocks((BlockTypes.Form,))
+    assert forms[0].html == corrected_html
+
+
+
+@pytest.mark.filename("table_ex2.pdf")
+@pytest.mark.config({"page_range": [0]})
+def test_llm_table_processor(pdf_document, detection_model, table_rec_model, recognition_model, mocker):
+    corrected_markdown = """
+| Column 1 | Column 2 | Column 3 | Column 4 |
+|----------|----------|----------|----------|
+| Value 1  | Value 2  | Value 3  | Value 4  |
+| Value 5  | Value 6  | Value 7  | Value 8  |
+| Value 9  | Value 10 | Value 11 | Value 12 |
+    """.strip()
+
+    mock_cls = Mock()
+    mock_cls.return_value.generate_response.return_value = {"corrected_markdown": corrected_markdown}
+    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+
+    cell_processor = TableProcessor(detection_model, recognition_model, table_rec_model)
+    cell_processor(pdf_document)
+
+    processor = LLMTableProcessor({"use_llm": True})
+    processor(pdf_document)
+
+    tables = pdf_document.contained_blocks((BlockTypes.Table,))
+    assert tables[0].cells[0].text == "Column 1"
+
+
+@pytest.mark.filename("adversarial.pdf")
+@pytest.mark.config({"page_range": [0]})
+def test_llm_text_processor(pdf_document, mocker):
+    inline_math_block = pdf_document.contained_blocks((BlockTypes.TextInlineMath,))[0]
+    text_lines = inline_math_block.contained_blocks(pdf_document, (BlockTypes.Line,))
+    corrected_lines = ["<i>Text</i>"] * len(text_lines)
+
+    mock_cls = Mock()
+    mock_cls.return_value.generate_response.return_value = {"corrected_lines": corrected_lines}
+    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+
+    processor = LLMTextProcessor({"use_llm": True})
+    processor(pdf_document)
+
+    contained_spans = text_lines[0].contained_blocks(pdf_document, (BlockTypes.Span,))
+    assert contained_spans[0].text == "Text\n" # Newline inserted at end of line
+    assert contained_spans[0].formats == ["italic"]