vlmparse 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

vlmparse/clients/docling.py

@@ -0,0 +1,146 @@
+import asyncio
+from io import BytesIO
+from typing import Literal
+
+import httpx
+from loguru import logger
+from PIL import Image
+from pydantic import Field
+
+from vlmparse.clients.pipe_utils.html_to_md_conversion import html_to_md_keep_tables
+from vlmparse.clients.pipe_utils.utils import clean_response
+from vlmparse.converter import BaseConverter, ConverterConfig
+from vlmparse.data_model.document import Page
+from vlmparse.servers.docker_server import DockerServerConfig
+
+
+class DoclingDockerServerConfig(DockerServerConfig):
+    """Configuration for Docling Serve using official image."""
+
+    model_name: str = "docling"
+    docker_image: str = Field(default="")
+    cpu_only: bool = False
+    command_args: list[str] = Field(default_factory=list)
+    server_ready_indicators: list[str] = Field(
+        default_factory=lambda: ["Application startup complete", "Uvicorn running"]
+    )
+    enable_ui: bool = False
+    docker_port: int = 5001
+    container_port: int = 5001
+    environment: dict[str, str] = Field(
+        default_factory=lambda: {
+            "DOCLING_SERVE_HOST": "0.0.0.0",
+            "DOCLING_SERVE_PORT": "5001",
+            "LOG_LEVEL": "DEBUG",  # Enable verbose logging
+            # Performance Tuning
+            # "UVICORN_WORKERS": "4",  # Increase web server workers (Default: 1)
+            # "DOCLING_SERVE_ENG_LOC_NUM_WORKERS": "4",  # Increase processing workers (Default: 2)
+            "DOCLING_NUM_THREADS": "32",  # Increase torch threads (Default: 4)
+        }
+    )
+
+    def model_post_init(self, __context):
+        """Set docker_image and gpu_device_ids based on cpu_only if not explicitly provided."""
+        if not self.docker_image:
+            if self.cpu_only:
+                self.docker_image = "quay.io/docling-project/docling-serve-cpu:latest"
+            else:
+                self.docker_image = "quay.io/docling-project/docling-serve:latest"
+
+        # For CPU-only mode, explicitly disable GPU by setting empty list
+        if self.cpu_only and self.gpu_device_ids is None:
+            self.gpu_device_ids = []
+
+        if self.enable_ui:
+            self.command_args.append("--enable-ui")
+
+    @property
+    def client_config(self):
+        return DoclingConverterConfig(base_url=f"http://localhost:{self.docker_port}")
+
+
+class DoclingConverterConfig(ConverterConfig):
+    """Configuration for Docling converter client."""
+
+    model_name: str = "docling"
+    base_url: str = "http://localhost:5001"
+    timeout: int = 300
+    api_kwargs: dict = {"output_format": "markdown", "image_export_mode": "referenced"}
+
+    def get_client(self, **kwargs) -> "DoclingConverter":
+        return DoclingConverter(config=self, **kwargs)
+
+
+def image_to_bytes(image: Image.Image) -> bytes:
+    # Convert image to bytes for file upload
+    img_byte_arr = BytesIO()
+    image.save(img_byte_arr, format="PNG")
+    img_bytes = img_byte_arr.getvalue()
+    return img_bytes
+
+
+class DoclingConverter(BaseConverter):
+    """Client for Docling Serve API using httpx."""
+
+    def __init__(
+        self,
+        config: DoclingConverterConfig,
+        num_concurrent_files: int = 10,
+        num_concurrent_pages: int = 10,
+        save_folder: str | None = None,
+        save_mode: Literal["document", "md", "md_page"] = "document",
+        debug: bool = False,
+        return_documents_in_batch_mode: bool = False,
+    ):
+        super().__init__(
+            config=config,
+            num_concurrent_files=num_concurrent_files,
+            num_concurrent_pages=num_concurrent_pages,
+            save_folder=save_folder,
+            save_mode=save_mode,
+            debug=debug,
+            return_documents_in_batch_mode=return_documents_in_batch_mode,
+        )
+
+    async def async_call_inside_page(self, page: Page) -> Page:
+        """Process a single page using Docling Serve API."""
+        img_bytes = await asyncio.to_thread(image_to_bytes, page.image)
+
+        data = self.config.api_kwargs
+        url = f"{self.config.base_url}/v1/convert/file"
+        logger.debug(f"Calling Docling API at: {url}")
+        files = {"files": ("image.png", img_bytes, "image/png")}
+
+        try:
+            async with httpx.AsyncClient(timeout=self.config.timeout) as client:
+                response = await client.post(
+                    url, files=files, data=data, headers={"Accept": "application/json"}
+                )
+            response.raise_for_status()
+
+            result = response.json()
+            logger.info(f"Docling API response status: {response.status_code}")
+
+            # Extract text from the response
+            # The response structure depends on the output format
+            if self.config.api_kwargs["output_format"] == "markdown":
+                text = result["document"]["md_content"]
+
+            elif self.config.api_kwargs["output_format"] == "text":
+                text = result["document"]["md_content"]
+
+            else:  # json or other formats
+                text = str(result)
+
+            logger.info(f"Extracted text length: {len(text)}")
+
+            # Clean and convert the response
+            text = clean_response(text)
+            text = html_to_md_keep_tables(text)
+            page.text = text
+
+        except Exception as e:
+            logger.error(f"Error processing page with Docling: {e}")
+            page.text = f"Error: {str(e)}"
+
+        return page

vlmparse/clients/dotsocr.py

@@ -0,0 +1,277 @@
+import json
+import math
+from pathlib import Path
+from typing import ClassVar, Literal
+
+from loguru import logger
+from PIL import Image
+from pydantic import Field
+
+from vlmparse.clients.openai_converter import (
+    OpenAIConverterClient,
+    OpenAIConverterConfig,
+)
+from vlmparse.clients.pipe_utils.html_to_md_conversion import html_to_md_keep_tables
+from vlmparse.clients.pipe_utils.utils import clean_response
+from vlmparse.data_model.document import BoundingBox, Item, Page
+from vlmparse.servers.docker_server import DEFAULT_MODEL_NAME, DockerServerConfig
+from vlmparse.utils import to_base64
+
+DOCKERFILE_DIR = Path(__file__).parent.parent.parent / "docker_pipelines"
+
+
+class DotsOCRDockerServerConfig(DockerServerConfig):
+    """Configuration for DotsOCR model."""
+
+    model_name: str = "rednote-hilab/dots.ocr"
+    docker_image: str = "dotsocr:latest"
+    dockerfile_dir: str = str(DOCKERFILE_DIR / "dotsocr")
+    command_args: list[str] = Field(
+        default_factory=lambda: [
+            "--tensor-parallel-size",
+            "1",
+            "--gpu-memory-utilization",
+            "0.8",
+            "--chat-template-content-format",
+            "string",
+            "--served-model-name",
+            DEFAULT_MODEL_NAME,
+            "--trust-remote-code",
+            # "--limit-mm-per-prompt",
+            # '{"image": 1}',
+            # "--no-enable-prefix-caching",
+            # "--max-model-len",
+            # "16384",
+        ]
+    )
+    add_model_key_to_server: bool = False
+    aliases: list[str] = Field(default_factory=lambda: ["dotsocr"])
+
+    @property
+    def client_config(self):
+        return DotsOCRConverterConfig(llm_params=self.llm_params)
+
+
+class DotsOCRConverterConfig(OpenAIConverterConfig):
+    model_name: str = "rednote-hilab/dots.ocr"
+    preprompt: str | None = ""
+    postprompt: str | None = None
+    completion_kwargs: dict | None = {
+        "temperature": 0.1,
+        "top_p": 1.0,
+        "max_completion_tokens": 16384,
+    }
+    aliases: list[str] = Field(default_factory=lambda: ["dotsocr"])
+    dpi: int = 200
+    prompt_mode: Literal["prompt_layout_all_en", "prompt_ocr"] = "prompt_ocr"
+
+    def get_client(self, **kwargs) -> "DotsOCRConverter":
+        return DotsOCRConverter(config=self, **kwargs)
+
+
+class DotsOCRConverter(OpenAIConverterClient):
+    """DotsOCR VLLM converter."""
+
+    # Constants
+    MIN_PIXELS: ClassVar[int] = 3136
+    MAX_PIXELS: ClassVar[int] = 11289600
+    IMAGE_FACTOR: ClassVar[int] = 28
+
+    # Prompts
+    PROMPTS: ClassVar[dict] = {
+        "prompt_layout_all_en": """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
+
+1. Bbox format: [x1, y1, x2, y2]
+
+2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].
+
+3. Text Extraction & Formatting Rules:
+    - Picture: For the 'Picture' category, the text field should be omitted.
+    - Formula: Format its text as LaTeX.
+    - Table: Format its text as HTML.
+    - All Others (Text, Title, etc.): Format their text as Markdown.
+
+4. Constraints:
+    - The output text must be the original text from the image, with no translation.
+    - All layout elements must be sorted according to human reading order.
+
+5. Final Output: The entire output must be a single JSON object.
+""",
+        "prompt_ocr": """Extract the text content from this image.""",
+    }
+
+    @staticmethod
+    def round_by_factor(number: int, factor: int) -> int:
+        """Returns the closest integer to 'number' that is divisible by 'factor'."""
+        return round(number / factor) * factor
+
+    @staticmethod
+    def ceil_by_factor(number: int, factor: int) -> int:
+        """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+        return math.ceil(number / factor) * factor
+
+    @staticmethod
+    def floor_by_factor(number: int, factor: int) -> int:
+        """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+        return math.floor(number / factor) * factor
+
+    def smart_resize(
+        self,
+        height: int,
+        width: int,
+        factor: int = 28,
+        min_pixels: int = 3136,
+        max_pixels: int = 11289600,
+    ):
+        """Rescales image dimensions to meet factor, pixel range, and aspect ratio constraints."""
+        if max(height, width) / min(height, width) > 200:
+            raise ValueError(
+                f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
+            )
+        h_bar = max(factor, self.round_by_factor(height, factor))
+        w_bar = max(factor, self.round_by_factor(width, factor))
+        if h_bar * w_bar > max_pixels:
+            beta = math.sqrt((height * width) / max_pixels)
+            h_bar = max(factor, self.floor_by_factor(height / beta, factor))
+            w_bar = max(factor, self.floor_by_factor(width / beta, factor))
+        elif h_bar * w_bar < min_pixels:
+            beta = math.sqrt(min_pixels / (height * width))
+            h_bar = self.ceil_by_factor(height * beta, factor)
+            w_bar = self.ceil_by_factor(width * beta, factor)
+            if h_bar * w_bar > max_pixels:
+                beta = math.sqrt((h_bar * w_bar) / max_pixels)
+                h_bar = max(factor, self.floor_by_factor(h_bar / beta, factor))
+                w_bar = max(factor, self.floor_by_factor(w_bar / beta, factor))
+        return h_bar, w_bar
+
+    def fetch_image(
+        self,
+        image,
+        min_pixels=None,
+        max_pixels=None,
+    ) -> Image.Image:
+        """Fetch and resize image."""
+        # Resize if needed
+        if min_pixels or max_pixels:
+            width, height = image.size
+            if not min_pixels:
+                min_pixels = self.MIN_PIXELS
+            if not max_pixels:
+                max_pixels = self.MAX_PIXELS
+            resized_height, resized_width = self.smart_resize(
+                height,
+                width,
+                factor=self.IMAGE_FACTOR,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+            assert resized_height > 0 and resized_width > 0
+            image = image.resize((resized_width, resized_height))
+
+        return image
+
+    def post_process_cells(
+        self,
+        origin_image: Image.Image,
+        cells: list,
+        input_width: int,
+        input_height: int,
+    ) -> list:
+        """Post-process cell bounding boxes to original image dimensions."""
+        if not cells or not isinstance(cells, list):
+            return cells
+
+        original_width, original_height = origin_image.size
+
+        scale_x = input_width / original_width
+        scale_y = input_height / original_height
+
+        cells_out = []
+        for cell in cells:
+            bbox = cell["bbox"]
+            bbox_resized = [
+                int(float(bbox[0]) / scale_x),
+                int(float(bbox[1]) / scale_y),
+                int(float(bbox[2]) / scale_x),
+                int(float(bbox[3]) / scale_y),
+            ]
+            cell_copy = cell.copy()
+            cell_copy["bbox"] = bbox_resized
+            cells_out.append(cell_copy)
+
+        return cells_out
+
+    async def _async_inference_with_vllm(self, image, prompt):
+        """Run async inference with VLLM."""
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{to_base64(image)}"
+                        },
+                    },
+                    {"type": "text", "text": f"<|img|><|imgpad|><|endofimg|>{prompt}"},
+                ],
+            }
+        ]
+
+        return await self._get_chat_completion(messages)
+
+    async def _parse_image_vllm(self, origin_image, prompt_mode="prompt_layout_all_en"):
+        """Parse image using VLLM inference."""
+
+        image = self.fetch_image(
+            origin_image, min_pixels=self.MIN_PIXELS, max_pixels=self.MAX_PIXELS
+        )
+        prompt = self.PROMPTS[prompt_mode]
+
+        response = await self._async_inference_with_vllm(image, prompt)
+
+        if prompt_mode in ["prompt_layout_all_en"]:
+            try:
+                cells = json.loads(response)
+                cells = self.post_process_cells(
+                    origin_image,
+                    cells,
+                    image.width,
+                    image.height,
+                )
+                return {}, cells, False
+            except Exception as e:
+                logger.warning(f"cells post process error: {e}, returning raw response")
+                return {}, response, True
+        else:
+            return {}, response, None
+
+    async def async_call_inside_page(self, page: Page) -> Page:
+        image = page.image
+
+        _, response, _ = await self._parse_image_vllm(
+            image, prompt_mode=self.config.prompt_mode
+        )
+        logger.info("Response: " + str(response))
+
+        items = None
+        if self.config.prompt_mode == "prompt_layout_all_en":
+            text = "\n\n".join([item.get("text", "") for item in response])
+
+            items = []
+            for item in response:
+                l, t, r, b = item["bbox"]
+                items.append(
+                    Item(
+                        text=item.get("text", ""),
+                        box=BoundingBox(l=l, t=t, r=r, b=b),
+                        category=item["category"],
+                    )
+                )
+            response = text
+            page.items = items
+
+        text = clean_response(response)
+        text = html_to_md_keep_tables(text)
+        page.text = text
+        return page

vlmparse/clients/granite_docling.py

@@ -0,0 +1,132 @@
+from pydantic import Field
+
+from vlmparse.clients.openai_converter import (
+    OpenAIConverterClient,
+    OpenAIConverterConfig,
+)
+from vlmparse.clients.pipe_utils.html_to_md_conversion import html_to_md_keep_tables
+from vlmparse.clients.pipe_utils.utils import clean_response
+from vlmparse.data_model.document import Page
+from vlmparse.servers.docker_server import VLLMDockerServerConfig
+from vlmparse.utils import to_base64
+
+
+class GraniteDoclingDockerServerConfig(VLLMDockerServerConfig):
+    """Configuration for Granite Docling model."""
+
+    model_name: str = "ibm-granite/granite-docling-258M"
+    command_args: list[str] = Field(
+        default_factory=lambda: [
+            "--revision",
+            "untied",
+            "--limit-mm-per-prompt",
+            '{"image": 1}',
+            "--trust-remote-code",
+        ]
+    )
+    aliases: list[str] = Field(default_factory=lambda: ["granite-docling"])
+
+    @property
+    def client_config(self):
+        return GraniteDoclingConverterConfig(llm_params=self.llm_params)
+
+
+class GraniteDoclingConverterConfig(OpenAIConverterConfig):
+    """Granite Docling converter configuration."""
+
+    preprompt: str | None = None
+    postprompt: str | None = "Convert this page to docling."
+    completion_kwargs: dict | None = {
+        "temperature": 0.0,
+        "max_tokens": 8000,
+        "extra_body": {
+            "skip_special_tokens": False,
+        },
+    }
+    aliases: list[str] = Field(default_factory=lambda: ["granite-docling"])
+
+    def get_client(self, **kwargs) -> "GraniteDoclingConverter":
+        return GraniteDoclingConverter(config=self, **kwargs)
+
+
+class GraniteDoclingConverter(OpenAIConverterClient):
+    """Client for Granite Docling model."""
+
+    async def async_call_inside_page(self, page: Page) -> Page:
+        image = page.image.convert("RGB")
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{to_base64(image)}"
+                        },
+                    },
+                    {"type": "text", "text": self.config.postprompt},
+                ],
+            }
+        ]
+
+        doctags = await self._get_chat_completion_adaptive(
+            messages, completion_kwargs=self.config.completion_kwargs
+        )
+        doctags = clean_response(doctags)
+
+        page.raw_response = doctags
+        page.text = _doctags_to_markdown(doctags, image)
+        return page
+
+    async def _get_chat_completion_adaptive(
+        self, messages: list[dict], completion_kwargs: dict | None
+    ) -> str:
+        """
+        vLLM enforces input+output <= model context length. If `max_tokens` is too
+        high (especially for multimodal prompts), retry with progressively smaller
+        `max_tokens`.
+        """
+        kwargs = (completion_kwargs or {}).copy()
+        max_tokens = kwargs.get("max_tokens") or kwargs.get("max_completion_tokens")
+
+        for _ in range(6):
+            try:
+                return await self._get_chat_completion(
+                    messages, completion_kwargs=kwargs
+                )
+            except Exception as e:
+                msg = str(e)
+                too_large = (
+                    "max_tokens" in msg
+                    and "maximum context length" in msg
+                    and "is too large" in msg
+                )
+                if not too_large or not isinstance(max_tokens, int):
+                    raise
+
+                max_tokens = max(256, int(max_tokens * 0.75))
+                if "max_tokens" in kwargs:
+                    kwargs["max_tokens"] = max_tokens
+                if "max_completion_tokens" in kwargs:
+                    kwargs["max_completion_tokens"] = max_tokens
+
+        return await self._get_chat_completion(messages, completion_kwargs=kwargs)
+
+
+def _doctags_to_markdown(doctags: str, image):
+    try:
+        from docling_core.types.doc import DoclingDocument
+        from docling_core.types.doc.document import DocTagsDocument
+    except Exception as e:  # pragma: no cover
+        raise RuntimeError(
+            "Missing optional dependency 'docling-core'. "
+            "Install it with: pip install 'vlmparse[docling_core]'"
+        ) from e
+
+    doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image])
+    doc = DoclingDocument.load_from_doctags(doctags_doc, document_name="Document")
+
+    html = doc.export_to_html()
+    html = clean_response(html)
+    md = html_to_md_keep_tables(html, remove_head=True)
+    return md

vlmparse/clients/hunyuanocr.py

@@ -0,0 +1,45 @@
+from pydantic import Field
+
+from vlmparse.clients.openai_converter import OpenAIConverterConfig
+from vlmparse.servers.docker_server import VLLMDockerServerConfig
+
+
+class HunyuanOCRDockerServerConfig(VLLMDockerServerConfig):
+    """Configuration for HunyuanOCR model."""
+
+    model_name: str = "tencent/HunyuanOCR"
+    command_args: list[str] = Field(
+        default_factory=lambda: [
+            "--limit-mm-per-prompt",
+            '{"image": 1}',
+            "--async-scheduling",
+            "--no-enable-prefix-caching",
+            "--mm-processor-cache-gb",
+            "0",
+            # Default argument in the hunyuan model, not sure why it is set this low.
+            "--gpu-memory-utilization",
+            "0.2",
+        ]
+    )
+    aliases: list[str] = Field(default_factory=lambda: ["hunyuanocr"])
+
+    @property
+    def client_config(self):
+        return HunyuanOCRConverterConfig(llm_params=self.llm_params)
+
+
+class HunyuanOCRConverterConfig(OpenAIConverterConfig):
+    """HunyuanOCR converter"""
+
+    model_name: str = "tencent/HunyuanOCR"
+    preprompt: str | None = ""
+    postprompt: str | None = (
+        "Extract all information from the main body of the document image and represent it in markdown format, ignoring headers and footers. Tables should be expressed in HTML format, formulas in the document should be represented using LaTeX format, and the parsing should be organized according to the reading order."
+    )
+    completion_kwargs: dict | None = {
+        "temperature": 0.0,
+        "extra_body": {"top_k": 1, "repetition_penalty": 1.0},
+    }
+    max_image_size: int | None = 1540
+    dpi: int = 200
+    aliases: list[str] = Field(default_factory=lambda: ["hunyuanocr"])

vlmparse/clients/lightonocr.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+
+from pydantic import Field
+
+from vlmparse.clients.openai_converter import OpenAIConverterConfig
+from vlmparse.servers.docker_server import VLLMDockerServerConfig
+
+DOCKERFILE_DIR = Path(__file__).parent.parent.parent / "docker_pipelines"
+
+
+class LightOnOCRDockerServerConfig(VLLMDockerServerConfig):
+    """Configuration for LightOnOCR model."""
+
+    model_name: str = "lightonai/LightOnOCR-1B-1025"
+    command_args: list[str] = Field(
+        default_factory=lambda: [
+            "--limit-mm-per-prompt",
+            '{"image": 1}',
+            "--mm-processor-cache-gb",
+            "0",
+            "--no-enable-prefix-caching",
+        ]
+    )
+    aliases: list[str] = Field(default_factory=lambda: ["lightonocr"])
+
+    @property
+    def client_config(self):
+        return LightOnOCRConverterConfig(llm_params=self.llm_params)
+
+
+class LightOnOCRConverterConfig(OpenAIConverterConfig):
+    """LightOnOCR converter - backward compatibility alias."""
+
+    model_name: str = "lightonai/LightOnOCR-1B-1025"
+    preprompt: str | None = None
+    postprompt: str | None = None
+    completion_kwargs: dict | None = {
+        "temperature": 0.2,
+        "max_tokens": 4096,
+        "top_p": 0.9,
+    }
+    dpi: int = 200
+    aliases: list[str] = Field(default_factory=lambda: ["lightonocr"])