vlmparse 0.1.7__py3-none-any.whl → 0.1.8__py3-none-any.whl

vlmparse/build_doc.py

@@ -10,12 +10,10 @@ from .constants import PDF_EXTENSION
 
 
 def convert_pdfium(file_path, dpi):
-    pdf = pdfium.PdfDocument(file_path)
     pil_images = []
-    for page in pdf:
-        pil_images.append(page.render(scale=dpi / 72).to_pil())
-
-    pdf.close()
+    with pdfium.PdfDocument(file_path) as pdf:
+        for page in pdf:
+            pil_images.append(page.render(scale=dpi / 72).to_pil())
     return pil_images
 
 
@@ -32,24 +30,29 @@ def convert_pdfium_to_images(file_path, dpi=175):
         ]
 
     except PIL.Image.DecompressionBombError as e:
-        logger.exception(f"Got problem size document with {file_path}")
+        logger.opt(exception=True).warning(
+            "Decompression bomb detected for {file_path}, reducing DPI",
+            file_path=str(file_path),
+        )
         cur_size, limit_size = map(int, re.findall(r"\d+", str(e)))
         factor = custom_ceil(cur_size / limit_size, precision=1)
-        logger.warning(
-            f"Try again by reducing DPI for doc {file_path} from {dpi} to {dpi//factor}"
+        new_dpi = dpi // factor
+        logger.info(
+            "Retrying {file_path} with reduced DPI: {old_dpi} -> {new_dpi}",
+            file_path=str(file_path),
+            old_dpi=dpi,
+            new_dpi=new_dpi,
         )
-        dpi = dpi // factor
-        images = convert_pdfium(file_path, dpi=dpi)
+        images = convert_pdfium(file_path, dpi=new_dpi)
 
     return images
 
 
 def convert_specific_page_to_image(file_path, page_number, dpi=175):
-    pdf = pdfium.PdfDocument(file_path)
-    page = pdf.get_page(page_number)
-    image = page.render(scale=dpi / 72).to_pil()
-    image = image.convert("L").convert("RGB") if image.mode != "RGB" else image
-    pdf.close()
+    with pdfium.PdfDocument(file_path) as pdf:
+        page = pdf.get_page(page_number)
+        image = page.render(scale=dpi / 72).to_pil()
+        image = image.convert("L").convert("RGB") if image.mode != "RGB" else image
     return image
 
 
@@ -68,9 +71,7 @@ def resize_image(image, max_image_size):
 
 def get_page_count(file_path):
     if Path(file_path).suffix.lower() == PDF_EXTENSION:
-        pdf = pdfium.PdfDocument(file_path)
-        count = len(pdf)
-        pdf.close()
-        return count
+        with pdfium.PdfDocument(file_path) as pdf:
+            return len(pdf)
     else:
         return 1

vlmparse/cli.py

@@ -49,10 +49,19 @@ class DParseCLI:
         uri: str | None = None,
         gpus: str | None = None,
         mode: Literal["document", "md", "md_page"] = "document",
+        conversion_mode: Literal[
+            "ocr",
+            "ocr_layout",
+            "table",
+            "image_description",
+            "formula",
+            "chart",
+        ] = "ocr",
         with_vllm_server: bool = False,
         concurrency: int = 10,
         dpi: int | None = None,
         debug: bool = False,
+        _return_documents: bool = False,
     ):
         """Parse PDF documents and save results.
 
@@ -64,6 +73,7 @@ class DParseCLI:
             uri: URI of the server, if not specified and the pipe is vllm, a local server will be deployed
             gpus: Comma-separated GPU device IDs (e.g., "0" or "0,1,2"). If not specified, all GPUs will be used.
             mode: Output mode - "document" (save as JSON zip), "md" (save as markdown file), "md_page" (save as folder of markdown pages)
+            conversion_mode: Conversion mode - "ocr" (plain), "ocr_layout" (OCR with layout), "table" (table-centric), "image_description" (describe the image), "formula" (formula extraction), "chart" (chart recognition)
             with_vllm_server: If True, a local VLLM server will be deployed if the model is not found in the registry. Note that if the model is in the registry and the uri is None, the server will be anyway deployed.
             dpi: DPI to use for the conversion. If not specified, the default DPI will be used.
             debug: If True, run in debug mode (single-threaded, no concurrency)
@@ -76,9 +86,15 @@ class DParseCLI:
             gpus=gpus,
             with_vllm_server=with_vllm_server,
             concurrency=concurrency,
+            return_documents=_return_documents,
         ) as converter_with_server:
             return converter_with_server.parse(
-                inputs=inputs, out_folder=out_folder, mode=mode, dpi=dpi, debug=debug
+                inputs=inputs,
+                out_folder=out_folder,
+                mode=mode,
+                conversion_mode=conversion_mode,
+                dpi=dpi,
+                debug=debug,
             )
 
     def list(self):

vlmparse/clients/chandra.py

@@ -1,6 +1,9 @@
+import json
 import math
 import time
+from dataclasses import asdict, dataclass
 
+from bs4 import BeautifulSoup
 from loguru import logger
 from PIL import Image
 from pydantic import Field
@@ -11,7 +14,8 @@ from vlmparse.clients.openai_converter import (
 )
 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.data_model.box import BoundingBox
+from vlmparse.data_model.document import Item, Page
 from vlmparse.servers.docker_server import VLLMDockerServerConfig
 from vlmparse.utils import to_base64
 
@@ -110,11 +114,6 @@ OCR this image to HTML.
 {PROMPT_ENDING}
 """.strip()
 
-PROMPT_MAPPING = {
-    "ocr_layout": OCR_LAYOUT_PROMPT,
-    "ocr": OCR_PROMPT,
-}
-
 
 def scale_to_fit(
     img: Image.Image,
@@ -188,11 +187,135 @@ def detect_repeat_token(
     return False
 
 
+@dataclass
+class LayoutBlock:
+    """Represents a layout block with bounding box and content."""
+
+    bbox: list[int]
+    label: str
+    content: str
+
+
+def parse_layout(
+    html: str, image: Image.Image, bbox_scale: int = 1024
+) -> list[LayoutBlock]:
+    """
+    Parse HTML layout blocks with bounding boxes.
+
+    Args:
+        html: HTML string with layout blocks (divs with data-bbox and data-label attributes)
+        image: PIL Image to get dimensions for bbox scaling
+        bbox_scale: The scale used in the prompt for normalized bboxes
+
+    Returns:
+        List of LayoutBlock objects with scaled bounding boxes
+    """
+    soup = BeautifulSoup(html, "html.parser")
+    top_level_divs = soup.find_all("div", recursive=False)
+    width, height = image.size
+    width_scaler = width / bbox_scale
+    height_scaler = height / bbox_scale
+    layout_blocks = []
+
+    for div in top_level_divs:
+        bbox = div.get("data-bbox")
+
+        try:
+            bbox = json.loads(bbox)
+            assert len(bbox) == 4, "Invalid bbox length"
+        except Exception:
+            try:
+                bbox = bbox.split(" ")
+                assert len(bbox) == 4, "Invalid bbox length"
+            except Exception:
+                # Default bbox if parsing fails
+                bbox = [0, 0, bbox_scale, bbox_scale]
+
+        bbox = list(map(int, bbox))
+        # Scale bbox to image dimensions
+        bbox = [
+            max(0, int(bbox[0] * width_scaler)),
+            max(0, int(bbox[1] * height_scaler)),
+            min(int(bbox[2] * width_scaler), width),
+            min(int(bbox[3] * height_scaler), height),
+        ]
+
+        label = div.get("data-label", "block")
+        content = str(div.decode_contents())
+        layout_blocks.append(LayoutBlock(bbox=bbox, label=label, content=content))
+
+    return layout_blocks
+
+
+def parse_chunks(html: str, image: Image.Image, bbox_scale: int = 1024) -> list[dict]:
+    """
+    Parse HTML layout blocks into dictionaries.
+
+    Args:
+        html: HTML string with layout blocks
+        image: PIL Image to get dimensions for bbox scaling
+        bbox_scale: The scale used in the prompt for normalized bboxes
+
+    Returns:
+        List of dictionaries with bbox, label, and content keys
+    """
+    layout = parse_layout(html, image, bbox_scale=bbox_scale)
+    chunks = [asdict(block) for block in layout]
+    return chunks
+
+
+def layout_blocks_to_items(
+    layout_blocks: list[LayoutBlock],
+) -> list[Item]:
+    """
+    Convert layout blocks to Item objects for the Page model.
+
+    Args:
+        layout_blocks: List of LayoutBlock objects
+
+    Returns:
+        List of Item objects with category, box, and text
+    """
+    items = []
+    for block in layout_blocks:
+        # Convert content HTML to markdown
+        try:
+            text = html_to_md_keep_tables(block.content)
+        except Exception as e:
+            logger.warning(f"Error converting block content to markdown: {e}")
+            text = block.content
+
+        # Create bounding box from [x0, y0, x1, y1] format
+        bbox = BoundingBox(
+            l=block.bbox[0],
+            t=block.bbox[1],
+            r=block.bbox[2],
+            b=block.bbox[3],
+        )
+
+        items.append(
+            Item(
+                category=block.label,
+                box=bbox,
+                text=text.strip(),
+            )
+        )
+
+    return items
+
+
 class ChandraConverterConfig(OpenAIConverterConfig):
     """Chandra converter configuration."""
 
     model_name: str = "datalab-to/chandra"
-    prompt_type: str = "ocr"  # Default prompt type
+    postprompt: str | None = None
+    prompts: dict[str, str] = {
+        "ocr": OCR_PROMPT,
+        "ocr_layout": OCR_LAYOUT_PROMPT,
+    }
+    prompt_mode_map: dict[str, str] = {
+        "table": "ocr_layout",
+    }
     bbox_scale: int = 1024
     max_retries: int = 0
     max_failure_retries: int = None
@@ -216,8 +339,7 @@ class ChandraConverterClient(OpenAIConverterClient):
 
     async def async_call_inside_page(self, page: Page) -> Page:
         """Process a single page using Chandra logic."""
-
-        prompt = PROMPT_MAPPING.get(self.config.prompt_type, OCR_PROMPT)
+        prompt = self.get_prompt_for_mode() or OCR_PROMPT
         prompt = prompt.replace("{bbox_scale}", str(self.config.bbox_scale))
 
         image = scale_to_fit(page.image)
@@ -238,61 +360,34 @@ class ChandraConverterClient(OpenAIConverterClient):
 
         retries = 0
         max_retries = self.config.max_retries
-        max_failure_retries = self.config.max_failure_retries
 
         result_content = ""
-        error_occurred = False
 
         while True:
-            try:
-                # Adjust temperature if retrying
-                temperature = self.config.completion_kwargs.get("temperature", 0.0)
-                if retries > 0:
-                    temperature = 0.3  # As per vllm.py logic
-
-                completion_kwargs = self.config.completion_kwargs.copy()
-                completion_kwargs["temperature"] = temperature
-                if retries > 0:
-                    completion_kwargs["top_p"] = 0.95
-
-                result_content = await self._get_chat_completion(
-                    messages, completion_kwargs=completion_kwargs
-                )
-                error_occurred = False
-            except Exception as e:
-                logger.error(f"Error during VLLM generation: {e}")
-                error_occurred = True
-                result_content = ""
-
             should_retry = False
-
-            # Check for repeat token
-            if not error_occurred:
-                has_repeat = detect_repeat_token(result_content) or (
-                    len(result_content) > 50
-                    and detect_repeat_token(result_content, cut_from_end=50)
+            # Adjust temperature if retrying
+            temperature = self.config.completion_kwargs.get("temperature", 0.0)
+            if retries > 0:
+                temperature = 0.3  # As per vllm.py logic
+
+            completion_kwargs = self.config.completion_kwargs.copy()
+            completion_kwargs["temperature"] = temperature
+            if retries > 0:
+                completion_kwargs["top_p"] = 0.95
+
+            result_content, usage = await self._get_chat_completion(
+                messages, completion_kwargs=completion_kwargs
+            )
+
+            has_repeat = detect_repeat_token(result_content) or (
+                len(result_content) > 50
+                and detect_repeat_token(result_content, cut_from_end=50)
+            )
+            if has_repeat and retries < max_retries:
+                logger.warning(
+                    f"Detected repeat token, retrying generation (attempt {retries + 1})..."
                 )
-                if has_repeat and retries < max_retries:
-                    logger.warning(
-                        f"Detected repeat token, retrying generation (attempt {retries + 1})..."
-                    )
-                    should_retry = True
-
-            # Check for error
-            if error_occurred:
-                if max_failure_retries is not None:
-                    if retries < max_failure_retries:
-                        logger.warning(
-                            f"Detected vllm error, retrying generation (attempt {retries + 1})..."
-                        )
-                        should_retry = True
-                elif (
-                    retries < max_retries
-                ):  # Fallback to max_retries if max_failure_retries not set (vllm.py logic varies slightly but this is safe)
-                    logger.warning(
-                        f"Detected vllm error, retrying generation (attempt {retries + 1})..."
-                    )
-                    should_retry = True
+                should_retry = True
 
             if should_retry:
                 time.sleep(2 * (retries + 1))
@@ -305,10 +400,27 @@ class ChandraConverterClient(OpenAIConverterClient):
         page.raw_response = result_content
         text = clean_response(result_content)
 
+        # Check if we're in layout mode (ocr_layout prompt)
+        current_prompt_key = self.get_prompt_key()
+        is_layout_mode = current_prompt_key == "ocr_layout"
+
+        if is_layout_mode:
+            # Parse layout blocks and populate items
+            try:
+                layout_blocks = parse_layout(
+                    text, image, bbox_scale=self.config.bbox_scale
+                )
+                page.items = layout_blocks_to_items(layout_blocks)
+                logger.info(f"Parsed {len(page.items)} layout blocks")
+            except Exception as e:
+                logger.warning(f"Error parsing layout blocks: {e}")
+                page.items = []
+
         # Convert HTML to MD
         text = html_to_md_keep_tables(text)
         page.text = text
-
+        page.completion_tokens = usage.completion_tokens
+        page.prompt_tokens = usage.prompt_tokens
         return page
 
 
@@ -320,4 +432,8 @@ class ChandraDockerServerConfig(VLLMDockerServerConfig):
 
     @property
     def client_config(self):
-        return ChandraConverterConfig(llm_params=self.llm_params)
+        return ChandraConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )

vlmparse/clients/deepseekocr.py

@@ -1,5 +1,4 @@
 import re
-from typing import ClassVar, Literal
 
 from loguru import logger
 from PIL import Image
@@ -35,7 +34,11 @@ class DeepSeekOCRDockerServerConfig(VLLMDockerServerConfig):
 
     @property
     def client_config(self):
-        return DeepSeekOCRConverterConfig(llm_params=self.llm_params)
+        return DeepSeekOCRConverterConfig(
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
+            )
+        )
 
 
 class DeepSeekOCRConverterConfig(OpenAIConverterConfig):
@@ -43,8 +46,17 @@ class DeepSeekOCRConverterConfig(OpenAIConverterConfig):
 
     model_name: str = "deepseek-ai/DeepSeek-OCR"
     aliases: list[str] = Field(default_factory=lambda: ["deepseekocr"])
+    postprompt: str | None = None
+    prompts: dict[str, str] = {
+        "layout": "<|grounding|>Convert the document to markdown.",
+        "ocr": "Free OCR.",
+        "image_description": "Describe this image in detail.",
+    }
+    prompt_mode_map: dict[str, str] = {
+        "ocr_layout": "layout",
+        "table": "layout",
+    }
 
-    prompt_mode: Literal["layout", "ocr"] = "ocr"
     completion_kwargs: dict | None = {
         "temperature": 0.0,
         "max_tokens": 8181,
@@ -95,12 +107,6 @@ def extract_coordinates_and_label(ref_text):
 class DeepSeekOCRConverterClient(OpenAIConverterClient):
     """Client for DeepSeekOCR with specific post-processing."""
 
-    PROMPTS: ClassVar[dict] = {
-        "layout": "<|grounding|>Convert the document to markdown.",
-        "ocr": "Free OCR.",
-        "image_description": "Describe this image in detail.",
-    }
-
     def extract_items(self, image: Image.Image, matches: list) -> list[Item]:
         items = []
         width, height = image.size
@@ -153,6 +159,8 @@ class DeepSeekOCRConverterClient(OpenAIConverterClient):
         # Prepare messages as in parent class
         image = page.image
 
+        prompt_key = self.get_prompt_key() or "ocr"
+
         messages = [
             {
                 "role": "user",
@@ -163,17 +171,17 @@ class DeepSeekOCRConverterClient(OpenAIConverterClient):
                             "url": f"data:image/png;base64,{to_base64(image)}"
                         },
                     },
-                    {"type": "text", "text": self.PROMPTS[self.config.prompt_mode]},
+                    {"type": "text", "text": self.config.prompts[prompt_key]},
                 ],
             },
         ]
 
         # Get raw response using parent's method
-        response = await self._get_chat_completion(messages)
+        response, usage = await self._get_chat_completion(messages)
         logger.info("Response length: " + str(len(response)))
         page.raw_response = response
 
-        if self.config.prompt_mode == "layout":
+        if prompt_key == "layout":
             # Post-processing
             matches, matches_image, matches_other = re_match(response)
 
@@ -199,5 +207,8 @@ class DeepSeekOCRConverterClient(OpenAIConverterClient):
 
         page.text = outputs.strip()
         logger.debug(page.text)
+        if usage is not None:
+            page.prompt_tokens = usage.prompt_tokens
+            page.completion_tokens = usage.completion_tokens
 
         return page

vlmparse/clients/docling.py

@@ -62,7 +62,6 @@ class DoclingDockerServerConfig(DockerServerConfig):
 class DoclingConverterConfig(ConverterConfig):
     """Configuration for Docling converter client."""
 
-    base_url: str
     model_name: str = "docling"
     timeout: int = 300
     api_kwargs: dict = {"output_format": "markdown", "image_export_mode": "referenced"}

vlmparse/clients/dotsocr.py

@@ -1,14 +1,13 @@
 import json
 import math
 from pathlib import Path
-from typing import ClassVar, Literal
+from typing import ClassVar
 
 from loguru import logger
 from PIL import Image
 from pydantic import Field
 
 from vlmparse.clients.openai_converter import (
-    LLMParams,
     OpenAIConverterClient,
     OpenAIConverterConfig,
 )
@@ -48,12 +47,13 @@ class DotsOCRDockerServerConfig(DockerServerConfig):
     )
     add_model_key_to_server: bool = True
     aliases: list[str] = Field(default_factory=lambda: ["dotsocr"])
+    default_model_name: str = DEFAULT_MODEL_NAME
 
     @property
     def client_config(self):
         return DotsOCRConverterConfig(
-            llm_params=LLMParams(
-                base_url=f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}",
+            **self._create_client_kwargs(
+                f"http://localhost:{self.docker_port}{self.get_base_url_suffix()}"
             )
         )
 
@@ -65,29 +65,7 @@ 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] = {
+    prompts: dict[str, str] = {
         "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]
@@ -108,6 +86,30 @@ class DotsOCRConverter(OpenAIConverterClient):
 """,
         "prompt_ocr": """Extract the text content from this image.""",
     }
+    prompt_mode_map: dict[str, str] = {
+        "ocr": "prompt_ocr",
+        "ocr_layout": "prompt_layout_all_en",
+        "table": "prompt_layout_all_en",
+    }
+    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
+
+    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
 
     @staticmethod
     def round_by_factor(number: int, factor: int) -> int:
@@ -235,7 +237,7 @@ class DotsOCRConverter(OpenAIConverterClient):
         image = self.fetch_image(
             origin_image, min_pixels=self.MIN_PIXELS, max_pixels=self.MAX_PIXELS
         )
-        prompt = self.PROMPTS[prompt_mode]
+        prompt = self.config.prompts[prompt_mode]
 
         response, usage = await self._async_inference_with_vllm(image, prompt)
 
@@ -258,13 +260,15 @@ class DotsOCRConverter(OpenAIConverterClient):
     async def async_call_inside_page(self, page: Page) -> Page:
         image = page.image
 
+        prompt_key = self.get_prompt_key() or "prompt_ocr"
+
         _, response, _, usage = await self._parse_image_vllm(
-            image, prompt_mode=self.config.prompt_mode
+            image, prompt_mode=prompt_key
         )
         logger.info("Response: " + str(response))
 
         items = None
-        if self.config.prompt_mode == "prompt_layout_all_en":
+        if prompt_key == "prompt_layout_all_en":
             text = "\n\n".join([item.get("text", "") for item in response])
 
             items = []
@@ -286,5 +290,4 @@ class DotsOCRConverter(OpenAIConverterClient):
 
         page.completion_tokens = usage.completion_tokens
         page.prompt_tokens = usage.prompt_tokens
-        page.reasoning_tokens = usage.reasoning_tokens
         return page