Refactor llm services

marker/processors/llm/llm_image_description.py

@@ -9,15 +9,15 @@ from typing import Annotated, List
 
 
 class LLMImageDescriptionProcessor(BaseLLMSimpleBlockProcessor):
-    block_types = (BlockTypes.Picture, BlockTypes.Figure,)
-    extract_images: Annotated[
-        bool,
-        "Extract images from the document."
-    ] = True
+    block_types = (
+        BlockTypes.Picture,
+        BlockTypes.Figure,
+    )
+    extract_images: Annotated[bool, "Extract images from the document."] = True
     image_description_prompt: Annotated[
         str,
         "The prompt to use for generating image descriptions.",
-        "Default is a string containing the Gemini prompt."
+        "Default is a string containing the Gemini prompt.",
     ] = """You are a document analysis expert who specializes in creating text descriptions for images.
 You will receive an image of a picture or figure.  Your job will be to create a short description of the image.
 **Instructions:**
@@ -41,26 +41,34 @@ In this figure, a bar chart titled "Fruit Preference Survey" is showing the numb
 
     def inference_blocks(self, document: Document) -> List[BlockData]:
         blocks = super().inference_blocks(document)
+        if self.extract_images:
+            return []
         return blocks
 
     def block_prompts(self, document: Document) -> List[PromptData]:
         prompt_data = []
         for block_data in self.inference_blocks(document):
             block = block_data["block"]
-            prompt = self.image_description_prompt.replace("{raw_text}", block.raw_text(document))
+            prompt = self.image_description_prompt.replace(
+                "{raw_text}", block.raw_text(document)
+            )
             image = self.extract_image(document, block)
 
-            prompt_data.append({
-                "prompt": prompt,
-                "image": image,
-                "block": block,
-                "schema": ImageSchema,
-                "page": block_data["page"]
-            })
+            prompt_data.append(
+                {
+                    "prompt": prompt,
+                    "image": image,
+                    "block": block,
+                    "schema": ImageSchema,
+                    "page": block_data["page"],
+                }
+            )
 
         return prompt_data
 
-    def rewrite_block(self, response: dict, prompt_data: PromptData, document: Document):
+    def rewrite_block(
+        self, response: dict, prompt_data: PromptData, document: Document
+    ):
         block = prompt_data["block"]
 
         if not response or "image_description" not in response:
@@ -74,5 +82,6 @@ In this figure, a bar chart titled "Fruit Preference Survey" is showing the numb
 
         block.description = image_description
 
+
 class ImageSchema(BaseModel):
     image_description: str

marker/services/__init__.py

@@ -1,10 +1,12 @@
 from typing import Optional, List, Annotated
+from io import BytesIO
 
 import PIL
 from pydantic import BaseModel
 
 from marker.schema.blocks import Block
 from marker.util import assign_config, verify_config_keys
+import base64
 
 
 class BaseService:
@@ -14,6 +16,11 @@ class BaseService:
     ] = 2
     retry_wait_time: Annotated[int, "The wait time between retries."] = 3
 
+    def img_to_base64(self, img: PIL.Image.Image):
+        image_bytes = BytesIO()
+        img.save(image_bytes, format="WEBP")
+        return base64.b64encode(image_bytes.getvalue()).decode("utf-8")
+
     def process_images(self, images: List[PIL.Image.Image]) -> list:
         raise NotImplementedError
 

marker/services/azure_openai.py

@@ -1,8 +1,6 @@
-import base64
 import json
 import time
-from io import BytesIO
-from typing import Annotated, List, Union
+from typing import Annotated, List
 
 import PIL
 from marker.logger import get_logger
@@ -18,30 +16,17 @@ logger = get_logger()
 
 class AzureOpenAIService(BaseService):
     azure_endpoint: Annotated[
-        str,
-        "The Azure OpenAI endpoint URL. No trailing slash."
+        str, "The Azure OpenAI endpoint URL. No trailing slash."
     ] = None
     azure_api_key: Annotated[
-        str,
-        "The API key to use for the Azure OpenAI service."
-    ] = None
-    azure_api_version: Annotated[
-        str,
-        "The Azure OpenAI API version to use."
+        str, "The API key to use for the Azure OpenAI service."
     ] = None
+    azure_api_version: Annotated[str, "The Azure OpenAI API version to use."] = None
     deployment_name: Annotated[
-        str,
-        "The deployment name for the Azure OpenAI model."
+        str, "The deployment name for the Azure OpenAI model."
     ] = None
 
-    def image_to_base64(self, image: PIL.Image.Image):
-        image_bytes = BytesIO()
-        image.save(image_bytes, format="WEBP")
-        return base64.b64encode(image_bytes.getvalue()).decode("utf-8")
-
-    def prepare_images(
-        self, images: Union[Image.Image, List[Image.Image]]
-    ) -> List[dict]:
+    def process_images(self, images: List[PIL.Image.Image]) -> list:
         if isinstance(images, Image.Image):
             images = [images]
 
@@ -49,10 +34,8 @@ class AzureOpenAIService(BaseService):
             {
                 "type": "image_url",
                 "image_url": {
-                    "url": "data:image/webp;base64,{}".format(
-                        self.image_to_base64(img)
-                    ),
-                }
+                    "url": "data:image/webp;base64,{}".format(self.img_to_base64(img)),
+                },
             }
             for img in images
         ]
@@ -60,8 +43,8 @@ class AzureOpenAIService(BaseService):
     def __call__(
         self,
         prompt: str,
-        image: PIL.Image.Image | List[PIL.Image.Image],
-        block: Block,
+        image: PIL.Image.Image | List[PIL.Image.Image] | None,
+        block: Block | None,
         response_schema: type[BaseModel],
         max_retries: int | None = None,
         timeout: int | None = None,
@@ -72,11 +55,8 @@ class AzureOpenAIService(BaseService):
         if timeout is None:
             timeout = self.timeout
 
-        if not isinstance(image, list):
-            image = [image]
-
         client = self.get_client()
-        image_data = self.prepare_images(image)
+        image_data = self.format_image_for_llm(image)
 
         messages = [
             {
@@ -94,7 +74,7 @@ class AzureOpenAIService(BaseService):
                 response = client.beta.chat.completions.parse(
                     extra_headers={
                         "X-Title": "Marker",
-                        "HTTP-Referer": "https://github.com/VikParuchuri/marker",
+                        "HTTP-Referer": "https://github.com/datalab-to/marker",
                     },
                     model=self.deployment_name,
                     messages=messages,
@@ -124,4 +104,4 @@ class AzureOpenAIService(BaseService):
             api_version=self.azure_api_version,
             azure_endpoint=self.azure_endpoint,
             api_key=self.azure_api_key,
-        ) 
+        )

marker/services/claude.py

@@ -1,7 +1,5 @@
-import base64
 import json
 import time
-from io import BytesIO
 from typing import List, Annotated, T
 
 import PIL
@@ -26,11 +24,6 @@ class ClaudeService(BaseService):
         int, "The maximum number of tokens to use for a single Claude request."
     ] = 8192
 
-    def img_to_base64(self, img: PIL.Image.Image):
-        image_bytes = BytesIO()
-        img.save(image_bytes, format="WEBP")
-        return base64.b64encode(image_bytes.getvalue()).decode("utf-8")
-
     def process_images(self, images: List[Image.Image]) -> List[dict]:
         return [
             {

marker/services/ollama.py

@@ -1,6 +1,4 @@
-import base64
 import json
-from io import BytesIO
 from typing import Annotated, List
 
 import PIL
@@ -22,13 +20,8 @@ class OllamaService(BaseService):
         "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 process_images(self, images):
-        image_bytes = [self.image_to_base64(img) for img in images]
+        image_bytes = [self.img_to_base64(img) for img in images]
         return image_bytes
 
     def __call__(

marker/services/openai.py

@@ -1,7 +1,5 @@
-import base64
 import json
 import time
-from io import BytesIO
 from typing import Annotated, List
 
 import openai
@@ -32,21 +30,6 @@ class OpenAIService(BaseService):
         "The image format to use for the OpenAI-like service. Use 'png' for better compatability",
     ] = "webp"
 
-    def image_to_base64(self, image: PIL.Image.Image) -> str:
-        """
-        Convert PIL image to base64 string
-
-        Args:
-            image: Input PIL Image
-            format: Format to use for the image; use "png" for better compatability.
-
-        Returns:
-            Base-64 encoded image (in PNG format) to pass to LLM.
-        """
-        image_bytes = BytesIO()
-        image.save(image_bytes, format=self.openai_image_format)
-        return base64.b64encode(image_bytes.getvalue()).decode("utf-8")
-
     def process_images(self, images: List[Image.Image]) -> List[dict]:
         """
         Generate the base-64 encoded message to send to an
@@ -67,7 +50,7 @@ class OpenAIService(BaseService):
                 "type": "image_url",
                 "image_url": {
                     "url": "data:image/{};base64,{}".format(
-                        self.openai_image_format, self.image_to_base64(img)
+                        self.openai_image_format, self.img_to_base64(img)
                     ),
                 },
             }