Factor out llm services, enable local models

README.md

@@ -22,7 +22,7 @@ See [below](#benchmarks) for detailed speed and accuracy benchmarks, and instruc
 
 ## Hybrid Mode
 
-For the highest accuracy, pass the `--use_llm` flag to use an LLM alongside marker.  This will do things like merge tables across pages, format tables properly, and extract values from forms.  It uses `gemini-flash-2.0`, which is cheap and fast.
+For the highest accuracy, pass the `--use_llm` flag to use an LLM alongside marker.  This will do things like merge tables across pages, handle inline math, format tables properly, and extract values from forms.  It can use any Google model (`gemini-2.0-flash` by default), or any ollama model.  See [below](#llm-services) for details.
 
 Here is a table benchmark comparing marker, gemini flash alone, and marker with use_llm:
 
@@ -42,7 +42,7 @@ As you can see, the use_llm mode offers higher accuracy than marker or gemini al
 
 I want marker to be as widely accessible as possible, while still funding my development/training costs.  Research and personal usage is always okay, but there are some restrictions on commercial usage.
 
-The weights for the models are licensed `cc-by-nc-sa-4.0`, but I will waive that for any organization under $5M USD in gross revenue in the most recent 12-month period AND under $5M in lifetime VC/angel funding raised. You also must not be competitive with the [Datalab API](https://www.datalab.to/).  If you want to remove the GPL license requirements (dual-license) and/or use the weights commercially over the revenue limit, check out the options [here](https://www.datalab.to).
+The weights for the models are licensed `cc-by-nc-sa-4.0`, but I will waive that for any organization under \$5M USD in gross revenue in the most recent 12-month period AND under $5M in lifetime VC/angel funding raised. You also must not be competitive with the [Datalab API](https://www.datalab.to/).  If you want to remove the GPL license requirements (dual-license) and/or use the weights commercially over the revenue limit, check out the options [here](https://www.datalab.to).
 
 # Hosted API
 
@@ -105,6 +105,8 @@ Options:
 - `--languages TEXT`: Optionally specify which languages to use for OCR processing. Accepts a comma-separated list. Example: `--languages "en,fr,de"` for English, French, and German.
 - `config --help`: List all available builders, processors, and converters, and their associated configuration.  These values can be used to build a JSON configuration file for additional tweaking of marker defaults.
 - `--converter_cls`: One of `marker.converters.pdf.PdfConverter` (default) or `marker.converters.table.TableConverter`.  The `PdfConverter` will convert the whole PDF, the `TableConverter` will only extract and convert tables.
+- `--llm_service`: Which llm service to use if `--use_llm` is passed.  This defaults to `marker.services.gemini.GoogleGeminiService`.
+- `--help`: see all of the flags that can be passed into marker.  (it supports many more options then are listed above)
 
 The list of supported languages for surya OCR is [here](https://github.com/VikParuchuri/surya/blob/master/surya/recognition/languages.py).  If you don't need OCR, marker can work with any language.
 
@@ -146,7 +148,7 @@ text, _, images = text_from_rendered(rendered)
 
 ### Custom configuration
 
-You can pass configuration using the `ConfigParser`:
+You can pass configuration using the `ConfigParser`.  To see all available options, do `marker_single --help`.
 
 ```python
 from marker.converters.pdf import PdfConverter
@@ -310,6 +312,14 @@ All output formats will return a metadata dictionary, with the following fields:
 }
 ```
 
+# LLM Services
+
+When running with the `--use_llm` flag, you have a choice of services you can use:
+
+- `Gemini` - this will use the Gemini developer API by default.  You'll need to pass `--gemini_api_key` to configuration.
+- `Google Vertex` - this will use vertex, which can be more reliable.  You'll need to pass `--vertex_project_id` and `--vertex_location`.  To use it, set `--llm_service=marker.services.vertex.GoogleVertexService`.
+- `Ollama` - this will use local models.  You can configure `--ollama_base_url` and `--ollama_model`. To use it, set `--llm_service=marker.services.vertex.OllamaService`.
+
 # Internals
 
 Marker is easy to extend.  The core units of marker are:

marker/builders/llm_layout.py

@@ -1,12 +1,12 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Annotated
+from typing import Annotated, Type
 
 from surya.layout import LayoutPredictor
 from tqdm import tqdm
 from pydantic import BaseModel
 
 from marker.builders.layout import LayoutBuilder
-from marker.services.google import GoogleModel
+from marker.services import BaseService
 from marker.providers.pdf import PdfProvider
 from marker.schema import BlockTypes
 from marker.schema.blocks import Block
@@ -97,10 +97,10 @@ Potential labels:
 Respond only with one of `Figure`, `Picture`, `ComplexRegion`, `Table`, or `Form`.
 """
 
-    def __init__(self, layout_model: LayoutPredictor, config=None):
+    def __init__(self, layout_model: LayoutPredictor, llm_service: BaseService, config=None):
         super().__init__(layout_model, config)
 
-        self.model = GoogleModel(self.google_api_key, self.model_name)
+        self.llm_service = llm_service
 
     def __call__(self, document: Document, provider: PdfProvider):
         super().__call__(document, provider)
@@ -158,7 +158,7 @@ Respond only with one of `Figure`, `Picture`, `ComplexRegion`, `Table`, or `Form
     def process_block_relabeling(self, document: Document, page: PageGroup, block: Block, prompt: str):
         image = self.extract_image(document, block)
 
-        response = self.model.generate_response(
+        response = self.llm_service(
             prompt,
             image,
             block,

marker/config/crawler.py

@@ -9,10 +9,11 @@ from marker.converters import BaseConverter
 from marker.processors import BaseProcessor
 from marker.providers import BaseProvider
 from marker.renderers import BaseRenderer
+from marker.services import BaseService
 
 
 class ConfigCrawler:
-    def __init__(self, base_classes=(BaseBuilder, BaseProcessor, BaseConverter, BaseProvider, BaseRenderer)):
+    def __init__(self, base_classes=(BaseBuilder, BaseProcessor, BaseConverter, BaseProvider, BaseRenderer, BaseService)):
         self.base_classes = base_classes
         self.class_config_map = {}
 

marker/config/parser.py

@@ -39,9 +39,9 @@ class ConfigParser:
         fn = click.option("--languages", type=str, default=None, help="Comma separated list of languages to use for OCR.")(fn)
 
         # we put common options here
-        fn = click.option("--google_api_key", type=str, default=None, help="Google API key for using LLMs.")(fn)
         fn = click.option("--use_llm", is_flag=True, default=False, help="Enable higher quality processing with LLMs.")(fn)
         fn = click.option("--converter_cls", type=str, default=None, help="Converter class to use.  Defaults to PDF converter.")(fn)
+        fn = click.option("--llm_service", type=str, default=None, help="LLM service to use - should be full import path, like marker.services.gemini.GoogleGeminiService")(fn)
 
         # enum options
         fn = click.option("--force_layout_block", type=click.Choice(choices=[t.name for t in BlockTypes]), default=None,)(fn)
@@ -74,8 +74,23 @@ class ConfigParser:
                 case _:
                     if k in crawler.attr_set:
                         config[k] = v
+
+        # Backward compatibility for google_api_key
+        if settings.GOOGLE_API_KEY:
+            config["gemini_api_key"] = settings.GOOGLE_API_KEY
+
         return config
 
+    def get_llm_service(self):
+        # Only return an LLM service when use_llm is enabled
+        if not self.cli_options["use_llm"]:
+            return None
+
+        service_cls = self.cli_options["llm_service"]
+        if service_cls is None:
+            service_cls = "marker.services.gemini.GoogleGeminiService"
+        return service_cls
+
     def get_renderer(self):
         match self.cli_options["output_format"]:
             case "json":
@@ -122,3 +137,4 @@ class ConfigParser:
     def get_base_filename(self, filepath: str):
         basename = os.path.basename(filepath)
         return os.path.splitext(basename)[0]
+

marker/converters/__init__.py

@@ -13,6 +13,7 @@ class BaseConverter:
     def __init__(self, config: Optional[BaseModel | dict] = None):
         assign_config(self, config)
         self.config = config
+        self.llm_service = None
 
     def __call__(self, *args, **kwargs):
         raise NotImplementedError
@@ -52,7 +53,8 @@ class BaseConverter:
 
         meta_processor = LLMSimpleBlockMetaProcessor(
             processor_lst=simple_llm_processors,
-            config=self.config
+            llm_service=self.llm_service,
+            config=self.config,
         )
         other_processors.insert(insert_position, meta_processor)
         return other_processors

marker/converters/pdf.py

@@ -1,4 +1,7 @@
 import os
+
+from marker.services.gemini import GoogleGeminiService
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"  # disables a tokenizers warning
 
 import inspect
@@ -86,7 +89,14 @@ class PdfConverter(BaseConverter):
         DebugProcessor,
     )
 
-    def __init__(self, artifact_dict: Dict[str, Any], processor_list: Optional[List[str]] = None, renderer: str | None = None, config=None):
+    def __init__(
+        self,
+        artifact_dict: Dict[str, Any],
+        processor_list: Optional[List[str]] = None,
+        renderer: str | None = None,
+        llm_service: str | None = None,
+        config=None
+    ):
         super().__init__(config)
 
         for block_type, override_block_type in self.override_map.items():
@@ -102,6 +112,14 @@ class PdfConverter(BaseConverter):
         else:
             renderer = MarkdownRenderer
 
+        if llm_service:
+            llm_service_cls = strings_to_classes([llm_service])[0]
+            llm_service = self.resolve_dependencies(llm_service_cls)
+
+        # Inject llm service into artifact_dict so it can be picked up by processors, etc.
+        artifact_dict["llm_service"] = llm_service
+        self.llm_service = llm_service
+
         self.artifact_dict = artifact_dict
         self.renderer = renderer
 

marker/processors/llm/__init__.py

@@ -8,11 +8,12 @@ from PIL import Image
 
 from marker.processors import BaseProcessor
 from marker.schema import BlockTypes
-from marker.services.google import GoogleModel
 from marker.schema.blocks import Block
 from marker.schema.document import Document
 from marker.schema.groups import PageGroup
+from marker.services import BaseService
 from marker.settings import settings
+from marker.util import assign_config
 
 
 class PromptData(TypedDict):
@@ -67,14 +68,14 @@ class BaseLLMProcessor(BaseProcessor):
     ] = False
     block_types = None
 
-    def __init__(self, config=None):
+    def __init__(self, llm_service: BaseService, config=None):
         super().__init__(config)
 
-        self.model = None
+        self.llm_service = None
         if not self.use_llm:
             return
 
-        self.model = GoogleModel(self.google_api_key, self.model_name)
+        self.llm_service = llm_service
 
     def extract_image(self, document: Document, image_block: Block, remove_blocks: Sequence[BlockTypes] | None = None) -> Image.Image:
         return image_block.get_image(
@@ -90,7 +91,7 @@ class BaseLLMComplexBlockProcessor(BaseLLMProcessor):
     A processor for using LLMs to convert blocks with more complex logic.
     """
     def __call__(self, document: Document):
-        if not self.use_llm or self.model is None:
+        if not self.use_llm or self.llm_service is None:
             return
 
         try:
@@ -125,6 +126,10 @@ class BaseLLMSimpleBlockProcessor(BaseLLMProcessor):
     A processor for using LLMs to convert single blocks.
     """
 
+    # Override init since we don't need an llmservice here
+    def __init__(self, config=None):
+        assign_config(self, config)
+
     def __call__(self, result: dict, prompt_data: PromptData, document: Document):
         try:
             self.rewrite_block(result, prompt_data, document)

marker/processors/llm/llm_meta.py

@@ -5,18 +5,19 @@ from tqdm import tqdm
 
 from marker.processors.llm import BaseLLMSimpleBlockProcessor, BaseLLMProcessor
 from marker.schema.document import Document
+from marker.services import BaseService
 
 
 class LLMSimpleBlockMetaProcessor(BaseLLMProcessor):
     """
     A wrapper for simple LLM processors, so they can all run in parallel.
     """
-    def __init__(self, processor_lst: List[BaseLLMSimpleBlockProcessor], config=None):
-        super().__init__(config)
+    def __init__(self, processor_lst: List[BaseLLMSimpleBlockProcessor], llm_service: BaseService, config=None):
+        super().__init__(llm_service, config)
         self.processors = processor_lst
 
     def __call__(self, document: Document):
-        if not self.use_llm or self.model is None:
+        if not self.use_llm or self.llm_service is None:
             return
 
         total = sum([len(processor.inference_blocks(document)) for processor in self.processors])
@@ -50,4 +51,4 @@ class LLMSimpleBlockMetaProcessor(BaseLLMProcessor):
         pbar.close()
 
     def get_response(self, prompt_data: Dict[str, Any]):
-        return self.model.generate_response(prompt_data["prompt"], prompt_data["image"], prompt_data["block"], prompt_data["schema"])
+        return self.llm_service(prompt_data["prompt"], prompt_data["image"], prompt_data["block"], prompt_data["schema"])

marker/processors/llm/llm_table.py

@@ -134,7 +134,7 @@ No corrections needed.
     def rewrite_single_chunk(self, page: PageGroup, block: Block, block_html: str, children: List[TableCell], image: Image.Image):
         prompt = self.table_rewriting_prompt.replace("{block_html}", block_html)
 
-        response = self.model.generate_response(prompt, image, block, TableSchema)
+        response = self.llm_service(prompt, image, block, TableSchema)
 
         if not response or "corrected_html" not in response:
             block.update_metadata(llm_error_count=1)

marker/processors/llm/llm_table_merge.py

@@ -240,7 +240,7 @@ Table 2
 
             prompt = self.table_merge_prompt.replace("{{table1}}", start_html).replace("{{table2}}", curr_html)
 
-            response = self.model.generate_response(
+            response = self.llm_service(
                 prompt,
                 [start_image, curr_image],
                 curr_block,

marker/scripts/convert.py

@@ -51,7 +51,8 @@ def process_single_pdf(args):
             config=config_parser.generate_config_dict(),
             artifact_dict=model_refs,
             processor_list=config_parser.get_processors(),
-            renderer=config_parser.get_renderer()
+            renderer=config_parser.get_renderer(),
+            llm_service=config_parser.get_llm_service()
         )
         rendered = converter(fpath)
         out_folder = config_parser.get_output_folder(fpath)

marker/scripts/convert_single.py

@@ -29,7 +29,8 @@ def convert_single_cli(fpath: str, **kwargs):
         config=config_parser.generate_config_dict(),
         artifact_dict=models,
         processor_list=config_parser.get_processors(),
-        renderer=config_parser.get_renderer()
+        renderer=config_parser.get_renderer(),
+        llm_service=config_parser.get_llm_service()
     )
     rendered = converter(fpath)
     out_folder = config_parser.get_output_folder(fpath)

marker/scripts/server.py

@@ -95,7 +95,8 @@ async def _convert_pdf(params: CommonParams):
             config=config_dict,
             artifact_dict=app_data["models"],
             processor_list=config_parser.get_processors(),
-            renderer=config_parser.get_renderer()
+            renderer=config_parser.get_renderer(),
+            llm_service=config_parser.get_llm_service()
         )
         rendered = converter(params.filepath)
         text, _, images = text_from_rendered(rendered)

marker/scripts/streamlit_app.py

@@ -56,7 +56,8 @@ def convert_pdf(fname: str, config_parser: ConfigParser) -> (str, Dict[str, Any]
         config=config_dict,
         artifact_dict=model_dict,
         processor_list=config_parser.get_processors(),
-        renderer=config_parser.get_renderer()
+        renderer=config_parser.get_renderer(),
+        llm_service=config_parser.get_llm_service()
     )
     return converter(fname)
 

marker/services/__init__.py

@@ -0,0 +1,26 @@
+from typing import Optional, List
+
+import PIL
+from pydantic import BaseModel
+
+from marker.schema.blocks import Block
+from marker.util import assign_config, verify_config_keys
+
+
+class BaseService:
+    def __init__(self, config: Optional[BaseModel | dict] = None):
+        assign_config(self, config)
+
+        # Ensure we have all necessary fields filled out (API keys, etc.)
+        verify_config_keys(self)
+
+    def __call__(
+        self,
+        prompt: str,
+        image: PIL.Image.Image | List[PIL.Image.Image],
+        block: Block,
+        response_schema: type[BaseModel],
+        max_retries: int = 1,
+        timeout: int = 15
+     ):
+        raise NotImplementedError

marker/services/{google.py → gemini.py}

@@ -1,7 +1,7 @@
 import json
 import time
 from io import BytesIO
-from typing import List
+from typing import List, Annotated
 
 import PIL
 from google import genai
@@ -10,29 +10,23 @@ from google.genai.errors import APIError
 from pydantic import BaseModel
 
 from marker.schema.blocks import Block
-from marker.settings import settings
+from marker.services import BaseService
 
-
-class GoogleModel:
-    def __init__(self, api_key: str, model_name: str):
-        if api_key is None:
-            raise ValueError("Google API key is not set")
-
-        self.api_key = api_key
-        self.model_name = model_name
-
-    def get_google_client(self, timeout: int = 60):
-        return genai.Client(
-            api_key=settings.GOOGLE_API_KEY,
-            http_options={"timeout": timeout * 1000} # Convert to milliseconds
-        )
+class BaseGeminiService(BaseService):
+    gemini_model_name: Annotated[
+        str,
+        "The name of the Google model to use for the service."
+    ] = "gemini-2.0-flash"
 
     def img_to_bytes(self, img: PIL.Image.Image):
         image_bytes = BytesIO()
         img.save(image_bytes, format="PNG")
         return image_bytes.getvalue()
 
-    def generate_response(
+    def get_google_client(self, timeout: int = 60):
+        raise NotImplementedError
+
+    def __call__(
             self,
             prompt: str,
             image: PIL.Image.Image | List[PIL.Image.Image],
@@ -51,7 +45,7 @@ class GoogleModel:
         while tries < max_retries:
             try:
                 responses = client.models.generate_content(
-                    model="gemini-2.0-flash",
+                    model=self.gemini_model_name,
                     contents=image_parts + [prompt], # According to gemini docs, it performs better if the image is the first element
                     config={
                         "temperature": 0,
@@ -78,3 +72,16 @@ class GoogleModel:
                 break
 
         return {}
+
+
+class GoogleGeminiService(BaseGeminiService):
+    gemini_api_key: Annotated[
+        str,
+        "The Google API key to use for the service."
+    ] = None
+
+    def get_google_client(self, timeout: int = 60):
+        return genai.Client(
+            api_key=self.gemini_api_key,
+            http_options={"timeout": timeout * 1000} # Convert to milliseconds
+        )

marker/services/ollama.py

@@ -0,0 +1,71 @@
+import base64
+import json
+from io import BytesIO
+from typing import Annotated, List
+
+import PIL
+import requests
+from pydantic import BaseModel
+
+from marker.schema.blocks import Block
+from marker.services import BaseService
+
+
+class OllamaService(BaseService):
+    ollama_base_url: Annotated[
+        str,
+        "The base url to use for ollama.  No trailing slash."
+    ] = "http://localhost:11434"
+    ollama_model: Annotated[
+        str,
+        "The model name to use for ollama."
+    ] = "llama3.2-vision"
+
+    def image_to_base64(self, image: PIL.Image.Image):
+        image_bytes = BytesIO()
+        image.save(image_bytes, format="PNG")
+        return base64.b64encode(image_bytes.getvalue()).decode("utf-8")
+
+    def __call__(
+        self,
+        prompt: str,
+        image: PIL.Image.Image | List[PIL.Image.Image],
+        block: Block,
+        response_schema: type[BaseModel],
+        max_retries: int = 1,
+        timeout: int = 15
+    ):
+        url = f"{self.ollama_base_url}/api/generate"
+        headers = {"Content-Type": "application/json"}
+
+        schema = response_schema.model_json_schema()
+        format_schema = {
+            "type": "object",
+            "properties": schema["properties"],
+            "required": schema["required"]
+        }
+
+        if not isinstance(image, list):
+            image = [image]
+
+        image_bytes = [self.image_to_base64(img) for img in image]
+
+        payload = {
+            "model": self.ollama_model,
+            "prompt": prompt,
+            "stream": False,
+            "format": format_schema,
+            "images": image_bytes
+        }
+
+        try:
+            response = requests.post(url, json=payload, headers=headers)
+            response.raise_for_status()
+            response_data = response.json()
+            data = response_data["response"]
+            print(data)
+            return json.loads(data)
+        except Exception as e:
+            print(f"Ollama inference failed: {e}")
+
+        return {}

marker/services/vertex.py

@@ -0,0 +1,23 @@
+from typing import Annotated
+
+from google import genai
+
+from marker.services.gemini import BaseGeminiService
+
+class GoogleVertexService(BaseGeminiService):
+    vertex_project_id: Annotated[
+        str,
+        "Google Cloud Project ID for Vertex AI.",
+    ] = None
+    vertex_location: Annotated[
+        str,
+        "Google Cloud Location for Vertex AI.",
+    ] = None
+
+    def get_google_client(self, timeout: int = 60):
+        return genai.Client(
+            vertexai=True,
+            project=self.vertex_project_id,
+            location=self.vertex_location,
+            http_options={"timeout": timeout * 1000} # Convert to milliseconds
+        )

marker/util.py

@@ -1,7 +1,6 @@
 import inspect
-import re
 from importlib import import_module
-from typing import List
+from typing import List, Annotated
 
 import numpy as np
 from pydantic import BaseModel
@@ -24,6 +23,19 @@ def classes_to_strings(items: List[type]) -> List[str]:
     return [f"{item.__module__}.{item.__name__}" for item in items]
 
 
+def verify_config_keys(obj):
+    annotations = inspect.get_annotations(obj.__class__)
+
+    none_vals = ""
+    for attr_name, annotation in annotations.items():
+        if isinstance(annotation, type(Annotated[str, ""])):
+            value = getattr(obj, attr_name)
+            if value is None:
+                none_vals += f"{attr_name}, "
+
+    assert len(none_vals) == 0, f"Missing values for {none_vals} are not allowed in {obj.__class__.__name__}."
+
+
 def assign_config(cls, config: BaseModel | dict | None):
     cls_name = cls.__class__.__name__
     if config is None:

tests/conftest.py

@@ -18,6 +18,7 @@ from marker.schema.blocks import Block
 from marker.renderers.markdown import MarkdownRenderer
 from marker.renderers.json import JSONRenderer
 from marker.schema.registry import register_block_class
+from marker.services.gemini import GoogleGeminiService
 from marker.util import classes_to_strings
 
 @pytest.fixture(scope="session")
@@ -126,6 +127,17 @@ def renderer(request, config):
     else:
         return MarkdownRenderer
 
+
+@pytest.fixture(scope="function")
+def llm_service(request):
+    llm_service = GoogleGeminiService(
+        config={
+            "gemini_api_key": "test"
+        }
+    )
+    yield llm_service
+
+
 @pytest.fixture(scope="function")
 def temp_image():
     img = Image.new("RGB", (512, 512), color="white")

tests/processors/test_inline_math.py

@@ -17,12 +17,11 @@ def test_llm_text_processor(pdf_document, mocker):
     corrected_lines = ["<math>Text</math>"] * 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)
+    mock_cls.return_value = {"corrected_lines": corrected_lines}
 
-    config = {"use_llm": True, "google_api_key": "test"}
+    config = {"use_llm": True, "gemini_api_key": "test"}
     processor_lst = [LLMTextProcessor(config)]
-    processor = LLMSimpleBlockMetaProcessor(processor_lst, config)
+    processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
 
     contained_spans = text_lines[0].contained_blocks(pdf_document, (BlockTypes.Span,))

tests/processors/test_llm_processors.py

@@ -14,11 +14,12 @@ from marker.renderers.markdown import MarkdownRenderer
 from marker.schema import BlockTypes
 from marker.schema.blocks import ComplexRegion
 
+
 @pytest.mark.filename("form_1040.pdf")
 @pytest.mark.config({"page_range": [0]})
-def test_llm_form_processor_no_config(pdf_document):
+def test_llm_form_processor_no_config(pdf_document, llm_service):
     processor_lst = [LLMFormProcessor()]
-    processor = LLMSimpleBlockMetaProcessor(processor_lst)
+    processor = LLMSimpleBlockMetaProcessor(processor_lst, llm_service)
     processor(pdf_document)
 
     forms = pdf_document.contained_blocks((BlockTypes.Form,))
@@ -27,9 +28,10 @@ def test_llm_form_processor_no_config(pdf_document):
 
 @pytest.mark.filename("form_1040.pdf")
 @pytest.mark.config({"page_range": [0]})
-def test_llm_form_processor_no_cells(pdf_document):
-    processor_lst = [LLMFormProcessor({"use_llm": True, "google_api_key": "test"})]
-    processor = LLMSimpleBlockMetaProcessor(processor_lst)
+def test_llm_form_processor_no_cells(pdf_document, llm_service):
+    config = {"use_llm": True, "google_api_key": "test"}
+    processor_lst = [LLMFormProcessor(config)]
+    processor = LLMSimpleBlockMetaProcessor(processor_lst, llm_service, config)
     processor(pdf_document)
 
     forms = pdf_document.contained_blocks((BlockTypes.Form,))
@@ -38,20 +40,19 @@ def test_llm_form_processor_no_cells(pdf_document):
 
 @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):
+def test_llm_form_processor(pdf_document, detection_model, table_rec_model, recognition_model):
     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_html": corrected_html}
-    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+    mock_cls.return_value = {"corrected_html": corrected_html}
 
     cell_processor = TableProcessor(detection_model, recognition_model, table_rec_model)
     cell_processor(pdf_document)
 
     config = {"use_llm": True, "google_api_key": "test"}
     processor_lst = [LLMFormProcessor(config)]
-    processor = LLMSimpleBlockMetaProcessor(processor_lst, config)
+    processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
 
     forms = pdf_document.contained_blocks((BlockTypes.Form,))
@@ -61,7 +62,7 @@ def test_llm_form_processor(pdf_document, detection_model, table_rec_model, reco
 
 @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):
+def test_llm_table_processor(pdf_document, detection_model, table_rec_model, recognition_model):
     corrected_html = """
 <table>
     <tr>
@@ -86,13 +87,12 @@ def test_llm_table_processor(pdf_document, detection_model, table_rec_model, rec
     """.strip()
 
     mock_cls = Mock()
-    mock_cls.return_value.generate_response.return_value = {"corrected_html": corrected_html}
-    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+    mock_cls.return_value = {"corrected_html": corrected_html}
 
     cell_processor = TableProcessor(detection_model, recognition_model, table_rec_model)
     cell_processor(pdf_document)
 
-    processor = LLMTableProcessor({"use_llm": True, "google_api_key": "test"})
+    processor = LLMTableProcessor(mock_cls, {"use_llm": True, "google_api_key": "test"})
     processor(pdf_document)
 
     tables = pdf_document.contained_blocks((BlockTypes.Table,))
@@ -107,8 +107,9 @@ def test_llm_table_processor(pdf_document, detection_model, table_rec_model, rec
 @pytest.mark.config({"page_range": [0]})
 def test_llm_caption_processor_disabled(pdf_document):
     config = {"use_llm": True, "google_api_key": "test"}
+    mock_cls = MagicMock()
     processor_lst = [LLMImageDescriptionProcessor(config)]
-    processor = LLMSimpleBlockMetaProcessor(processor_lst, config)
+    processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
 
     contained_pictures = pdf_document.contained_blocks((BlockTypes.Picture, BlockTypes.Figure))
@@ -116,15 +117,14 @@ def test_llm_caption_processor_disabled(pdf_document):
 
 @pytest.mark.filename("A17_FlightPlan.pdf")
 @pytest.mark.config({"page_range": [0]})
-def test_llm_caption_processor(pdf_document, mocker):
+def test_llm_caption_processor(pdf_document):
     description = "This is an image description."
     mock_cls = Mock()
-    mock_cls.return_value.generate_response.return_value = {"image_description": description}
-    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+    mock_cls.return_value = {"image_description": description}
 
     config = {"use_llm": True, "google_api_key": "test", "extract_images": False}
     processor_lst = [LLMImageDescriptionProcessor(config)]
-    processor = LLMSimpleBlockMetaProcessor(processor_lst, config)
+    processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
 
     contained_pictures = pdf_document.contained_blocks((BlockTypes.Picture, BlockTypes.Figure))
@@ -139,11 +139,10 @@ def test_llm_caption_processor(pdf_document, mocker):
 
 @pytest.mark.filename("A17_FlightPlan.pdf")
 @pytest.mark.config({"page_range": [0]})
-def test_llm_complex_region_processor(pdf_document, mocker):
+def test_llm_complex_region_processor(pdf_document):
     md = "This is some *markdown* for a complex region."
     mock_cls = Mock()
-    mock_cls.return_value.generate_response.return_value = {"corrected_markdown": md * 25}
-    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+    mock_cls.return_value = {"corrected_markdown": md * 25}
 
     # Replace the block with a complex region
     old_block = pdf_document.pages[0].children[0]
@@ -155,7 +154,7 @@ def test_llm_complex_region_processor(pdf_document, mocker):
     # Test processor
     config = {"use_llm": True, "google_api_key": "test"}
     processor_lst = [LLMComplexRegionProcessor(config)]
-    processor = LLMSimpleBlockMetaProcessor(processor_lst, config)
+    processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
 
     # Ensure the rendering includes the description
@@ -166,15 +165,14 @@ def test_llm_complex_region_processor(pdf_document, mocker):
 
 @pytest.mark.filename("adversarial.pdf")
 @pytest.mark.config({"page_range": [0]})
-def test_multi_llm_processors(pdf_document, mocker):
+def test_multi_llm_processors(pdf_document):
     description = "<math>This is an image description.  And here is a lot of writing about it.</math>" * 10
     mock_cls = Mock()
-    mock_cls.return_value.generate_response.return_value = {"image_description": description, "html_equation": description}
-    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
+    mock_cls.return_value = {"image_description": description, "html_equation": description}
 
     config = {"use_llm": True, "google_api_key": "test", "extract_images": False, "min_equation_height": .001}
     processor_lst = [LLMImageDescriptionProcessor(config), LLMEquationProcessor(config)]
-    processor = LLMSimpleBlockMetaProcessor(processor_lst, config)
+    processor = LLMSimpleBlockMetaProcessor(processor_lst, mock_cls, config)
     processor(pdf_document)
 
     contained_pictures = pdf_document.contained_blocks((BlockTypes.Picture, BlockTypes.Figure))

tests/processors/test_table_merge.py

@@ -10,11 +10,10 @@ from marker.schema import BlockTypes
 @pytest.mark.filename("table_ex2.pdf")
 def test_llm_table_processor_nomerge(pdf_document, detection_model, table_rec_model, recognition_model, mocker):
     mock_cls = Mock()
-    mock_cls.return_value.generate_response.return_value = {
+    mock_cls.return_value = {
         "merge": "true",
         "direction": "right"
     }
-    mocker.patch("marker.processors.llm.GoogleModel", mock_cls)
 
     cell_processor = TableProcessor(detection_model, recognition_model, table_rec_model)
     cell_processor(pdf_document)
@@ -22,7 +21,7 @@ def test_llm_table_processor_nomerge(pdf_document, detection_model, table_rec_mo
     tables = pdf_document.contained_blocks((BlockTypes.Table,))
     assert len(tables) == 3
 
-    processor = LLMTableMergeProcessor({"use_llm": True, "google_api_key": "test"})
+    processor = LLMTableMergeProcessor(mock_cls, {"use_llm": True, "google_api_key": "test"})
     processor(pdf_document)
 
     tables = pdf_document.contained_blocks((BlockTypes.Table,))