Merge pull request #631 from VikParuchuri/tarun-dev

marker/processors/table.py

@@ -1,4 +1,5 @@
 import re
+import html
 from collections import defaultdict
 from copy import deepcopy
 from typing import Annotated, List
@@ -158,7 +159,8 @@ class TableProcessor(BaseProcessor):
                 continue
             text = re.sub(r"(\s\.){2,}", "", text)  # Replace . . .
             text = re.sub(r"\.{2,}", "", text)  # Replace ..., like in table of contents
-            fixed_text.append(self.normalize_spaces(fix_text(text)))
+            text = self.normalize_spaces(fix_text(text))
+            fixed_text.append(html.escape(text))
         return fixed_text
 
     @staticmethod

marker/providers/pdf.py

@@ -247,7 +247,10 @@ class PdfProvider(BaseProvider):
                     )
             if self.check_line_spans(lines):
                 page_lines[page_id] = lines
-            self.page_refs[page_id] = page["refs"]
+            
+            self.page_refs[page_id] = []
+            if page_refs:= page.get('refs', None):
+                self.page_refs[page_id] = page_refs
 
         return page_lines
 

marker/services/openai.py

@@ -0,0 +1,117 @@
+import base64
+import json
+import time
+from io import BytesIO
+from typing import Annotated, List, Union
+
+import openai
+import PIL
+from openai import APITimeoutError, RateLimitError
+from PIL import Image
+from pydantic import BaseModel
+
+from marker.schema.blocks import Block
+from marker.services import BaseService
+
+
+class OpenAIService(BaseService):
+    openai_base_url: Annotated[
+        str,
+        "The base url to use for OpenAI-like models.  No trailing slash."
+    ] = "https://api.openai.com/v1"
+    openai_model: Annotated[
+        str,
+        "The model name to use for OpenAI-like model."
+    ] = "gpt-4o-mini"
+    openai_api_key: Annotated[
+        str,
+        "The API key to use for the OpenAI-like service."
+    ] = 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]:
+        if isinstance(images, Image.Image):
+            images = [images]
+
+        return [
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": "data:image/webp;base64,{}".format(
+                        self.image_to_base64(img)
+                    ),
+                }
+            }
+            for img in images
+        ]
+
+    def __call__(
+        self,
+        prompt: str,
+        image: PIL.Image.Image | List[PIL.Image.Image],
+        block: Block,
+        response_schema: type[BaseModel],
+        max_retries: int | None = None,
+        timeout: int | None = None,
+    ):
+        if max_retries is None:
+            max_retries = self.max_retries
+
+        if timeout is None:
+            timeout = self.timeout
+
+        if not isinstance(image, list):
+            image = [image]
+
+        client = self.get_client()
+        image_data = self.prepare_images(image)
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    *image_data,
+                    {"type": "text", "text": prompt},
+                ],
+            }
+        ]
+
+        tries = 0
+        while tries < max_retries:
+            try:
+                response = client.beta.chat.completions.parse(
+                    extra_headers={
+                        "X-Title": "Marker",
+                        "HTTP-Referer": "https://github.com/VikParuchuri/marker",
+                    },
+                    model=self.openai_model,
+                    messages=messages,
+                    timeout=timeout,
+                    response_format=response_schema,
+                )
+                response_text = response.choices[0].message.content
+                total_tokens = response.usage.total_tokens
+                block.update_metadata(llm_tokens_used=total_tokens, llm_request_count=1)
+                return json.loads(response_text)
+            except (APITimeoutError, RateLimitError) as e:
+                # Rate limit exceeded
+                tries += 1
+                wait_time = tries * 3
+                print(
+                    f"Rate limit error: {e}. Retrying in {wait_time} seconds... (Attempt {tries}/{max_retries})"
+                )
+                time.sleep(wait_time)
+            except Exception as e:
+                print(e)
+                break
+
+        return {}
+
+    def get_client(self) -> openai.OpenAI:
+        return openai.OpenAI(api_key=self.openai_api_key, base_url=self.openai_base_url)

pyproject.toml

@@ -43,6 +43,7 @@ openpyxl = {version = "^3.1.5", optional = true}
 python-pptx = {version = "^1.0.2", optional = true}
 ebooklib = {version = "^0.18", optional = true}
 weasyprint = {version = "^63.1", optional = true}
+openai = "^1.65.2"
 
 [tool.poetry.group.dev.dependencies]
 jupyter = "^1.0.0"

signatures/version1/cla.json

@@ -183,6 +183,14 @@
       "created_at": "2025-02-16T23:02:34Z",
       "repoId": 712111618,
       "pullRequestNo": 555
+    },
+    {
+      "name": "vicenciomf2",
+      "id": 127889973,
+      "comment_id": 2676007412,
+      "created_at": "2025-02-22T04:34:27Z",
+      "repoId": 712111618,
+      "pullRequestNo": 574
     }
   ]
 }