infinity-parser2 0.1.0__py3-none-any.whl

infinity_parser2/__init__.py

@@ -0,0 +1,28 @@
+"""Infinity-Parser2: Document parsing Python package."""
+
+__version__ = "0.1.0"
+
+from .parser import InfinityParser2
+from .backends import (
+    BaseBackend,
+    TransformersBackend,
+    VLLMEngineBackend,
+    VLLMServerBackend,
+)
+from .prompts import PROMPT_DOC2JSON, PROMPT_DOC2MD, SUPPORTED_TASK_TYPES
+from .utils import convert_pdf_to_images
+from .cli import main as cli_main
+
+__all__ = [
+    "InfinityParser2",
+    "BaseBackend",
+    "TransformersBackend",
+    "VLLMEngineBackend",
+    "VLLMServerBackend",
+    "convert_pdf_to_images",
+    "PROMPT_DOC2JSON",
+    "PROMPT_DOC2MD",
+    "SUPPORTED_TASK_TYPES",
+    "__version__",
+    "cli_main",
+]

infinity_parser2/__main__.py

@@ -0,0 +1,6 @@
+"""Entry point for running `python -m infinity_parser2`."""
+
+from .cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

infinity_parser2/backends/__init__.py

@@ -0,0 +1,13 @@
+"""Inference backends for Infinity-Parser2."""
+
+from .base import BaseBackend
+from .transformers import TransformersBackend
+from .vllm_engine import VLLMEngineBackend
+from .vllm_server import VLLMServerBackend
+
+__all__ = [
+    "BaseBackend",
+    "TransformersBackend",
+    "VLLMEngineBackend",
+    "VLLMServerBackend",
+]

infinity_parser2/backends/base.py

@@ -0,0 +1,61 @@
+"""Base backend interface for Infinity-Parser2."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Union
+
+from PIL import Image
+
+
+class BaseBackend(ABC):
+    """Abstract base class for inference backends.
+
+    All backends must implement init() and parse_batch() methods.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "infly/Infinity-Parser2-Pro",
+        device: str = "cuda",
+        **kwargs,
+    ):
+        """Initialize backend.
+
+        Args:
+            model_name: Model name on HuggingFace Hub or local path to the model.
+            device: Device type, "cuda" or "cpu".
+            **kwargs: Additional backend-specific arguments.
+        """
+        self.model_name = model_name
+        self.device = device
+        self.kwargs = kwargs
+
+    @abstractmethod
+    def init(self) -> None:
+        """Initialize the model and processor.
+
+        This method should be called before parse_batch().
+        """
+        pass
+
+    @abstractmethod
+    def parse_batch(
+        self,
+        input_data: list[Union[str, Image.Image]],
+        prompt: str,
+        batch_size: int = 1,
+        **kwargs,
+    ) -> list[str]:
+        """Parse multiple documents.
+
+        Args:
+            input_data: List of inputs, each can be:
+                - str: File path (image or PDF)
+                - PIL.Image.Image: Image object
+            prompt: Prompt text for the model.
+            batch_size: Maximum number of images to process in one batch.
+            **kwargs: Additional arguments passed to the model.
+
+        Returns:
+            List of parsed text content (one per input in the same order).
+        """
+        pass

infinity_parser2/backends/transformers.py

@@ -0,0 +1,159 @@
+"""Transformers backend for Infinity-Parser2."""
+
+import sys
+from typing import Union
+
+from PIL import Image
+import torch
+from tqdm import tqdm
+from transformers import AutoModelForImageTextToText, AutoProcessor
+from qwen_vl_utils import process_vision_info
+
+from ..utils import load_image
+from .base import BaseBackend
+
+
+class TransformersBackend(BaseBackend):
+    """Inference backend using HuggingFace transformers.
+
+    Supports local model inference with automatic device mapping.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "infly/Infinity-Parser2-Pro",
+        device: str = "cuda",
+        torch_dtype: str = "bfloat16",
+        min_pixels: int = 2048,
+        max_pixels: int = 16777216,
+        **kwargs,
+    ):
+        """Initialize Transformers backend.
+
+        Args:
+            model_name: Model name on HuggingFace Hub or local path.
+            device: Device type, "cuda" or "cpu".
+            torch_dtype: Data type for model weights, "float16" or "bfloat16".
+            min_pixels: Minimum number of pixels for image input.
+            max_pixels: Maximum number of pixels for image input.
+            **kwargs: Additional arguments for AutoModelForImageTextToText.from_pretrained.
+        """
+        super().__init__(model_name, device, **kwargs)
+        self.torch_dtype = getattr(torch, torch_dtype, torch.bfloat16)
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.init()
+
+    def init(self) -> None:
+        """Initialize the model and processor."""
+        self._model = AutoModelForImageTextToText.from_pretrained(
+            self.model_name,
+            torch_dtype=self.torch_dtype,
+            device_map="auto",
+            **self.kwargs,
+        )
+        self._processor = AutoProcessor.from_pretrained(self.model_name)
+
+    def _process_inputs(
+        self,
+        inputs: list[Union[str, Image.Image]],
+        prompt: str,
+        **kwargs,
+    ) -> dict:
+        """Process inputs for generation.
+
+        Returns:
+            Dictionary with processed inputs for the model.
+        """
+        images = [load_image(item) for item in inputs]
+
+        messages = [
+            [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": img, "min_pixels": self.min_pixels, "max_pixels": self.max_pixels},
+                        {"type": "text", "text": prompt},
+                    ],
+                }
+            ]
+            for img in images
+        ]
+
+        chat_template_kwargs = {"enable_thinking": False}
+
+        text = self._processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, **chat_template_kwargs
+        )
+        image_inputs, _ = process_vision_info(messages, image_patch_size=16)
+
+        inputs = self._processor(
+            text=text,
+            images=image_inputs,
+            do_resize=False,
+            padding=True,
+            return_tensors="pt",
+        )
+        inputs.pop("token_type_ids", None)
+        return inputs
+
+    def _generate(self, inputs: dict, **kwargs) -> list[str]:
+        """Run model generation and decode outputs.
+
+        Args:
+            inputs: Processed inputs from _process_inputs.
+            **kwargs: Generation arguments.
+
+        Returns:
+            List of generated text outputs.
+        """
+        # Move tensors to device
+        inputs = {
+            k: v.to(self._model.device) if isinstance(v, torch.Tensor) else v
+            for k, v in inputs.items()
+        }
+
+        generated_ids = self._model.generate(
+            **inputs,
+            max_new_tokens=kwargs.get("max_new_tokens", 32768),
+            temperature=kwargs.get("temperature", 0.0),
+            top_p=kwargs.get("top_p", 1.0),
+        )
+
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):]
+            for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
+        ]
+        output_text = self._processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+        return output_text
+
+    def parse_batch(
+        self,
+        input_data: list[Union[str, Image.Image]],
+        prompt: str,
+        batch_size: int = 1,
+        **kwargs,
+    ) -> list[str]:
+        """Parse multiple documents with batched inference.
+
+        Args:
+            input_data: List of file paths or PIL Images.
+            prompt: Prompt text for the model.
+            batch_size: Maximum number of images to process in one batch.
+            **kwargs: Additional arguments.
+
+        Returns:
+            List of parsed text content (one per input in the same order).
+        """
+        results = [None] * len(input_data)
+
+        for i in tqdm(range(0, len(input_data), batch_size), desc="Parsing", file=sys.stdout):
+            batch = input_data[i : i + batch_size]
+            inputs = self._process_inputs(batch, prompt, **kwargs)
+            batch_results = self._generate(inputs, **kwargs)
+            for j, result in enumerate(batch_results):
+                results[i + j] = result
+
+        return results

infinity_parser2/backends/vllm_engine.py

@@ -0,0 +1,117 @@
+"""vLLM Engine backend for Infinity-Parser2.
+
+Uses vLLM's LLM class for offline batch inference.
+"""
+
+import sys
+from typing import Union
+
+from PIL import Image
+from tqdm import tqdm
+from vllm import LLM, SamplingParams
+
+from .base import BaseBackend
+from ..utils import encode_file_to_base64
+
+
+class VLLMEngineBackend(BaseBackend):
+    """Offline inference backend using vLLM Engine.
+
+    Uses vLLM's LLM class for local batch inference with tensor parallelism.
+    Reference: https://docs.vllm.ai/en/latest/serving/offline_inference.html
+    """
+
+    def __init__(
+        self,
+        model_name: str = "infly/Infinity-Parser2-Pro",
+        device: str = "cuda",
+        tensor_parallel_size: int = 1,
+        min_pixels: int = 2048,
+        max_pixels: int = 16777216,
+        **kwargs,
+    ):
+        """Initialize vLLM Engine backend.
+
+        Args:
+            model_name: Model name on HuggingFace Hub or local path.
+            device: Device type, "cuda" or "cpu".
+            tensor_parallel_size: Number of GPUs for tensor parallelism.
+            **kwargs: Additional arguments for vllm.LLM.
+        """
+        super().__init__(model_name, device, **kwargs)
+        self.tensor_parallel_size = tensor_parallel_size
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.init()
+
+    def init(self) -> None:
+        """Initialize the vLLM LLM instance."""
+        # model_name can be a HuggingFace model ID or local path
+        self._llm = LLM(
+            model=self.model_name,
+            trust_remote_code=True,
+            tensor_parallel_size=self.tensor_parallel_size,
+            gpu_memory_utilization=0.85,
+            **self.kwargs,
+        )
+
+    def _build_messages(self, base64_data: str, mime_type: str, prompt: str) -> list:
+        """Build chat messages with base64-encoded image.
+
+        Returns:
+            List of message dictionaries.
+        """
+        return [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_data}"}},
+                    {"type": "text", "text": prompt},
+                ],
+            }
+        ]
+
+    def parse_batch(
+        self,
+        input_data: list[Union[str, Image.Image]],
+        prompt: str,
+        batch_size: int = 1,
+        **kwargs,
+    ) -> list[str]:
+        """Parse multiple documents with batched inference.
+
+        Args:
+            input_data: List of file paths or PIL Images.
+            prompt: Prompt text for the model.
+            batch_size: Maximum number of images to process in one batch.
+            **kwargs: Additional arguments.
+
+        Returns:
+            List of parsed text content (one per input in the same order).
+        """
+
+        sampling_params = SamplingParams(
+            max_tokens=kwargs.get("max_new_tokens", 32768),
+            temperature=kwargs.get("temperature", 0.0),
+            top_p=kwargs.get("top_p", 1.0),
+        )
+        chat_template_kwargs = {"enable_thinking": False}
+
+        all_messages = []
+        for item in input_data:
+            base64_data, mime_type = encode_file_to_base64(item, min_pixels=self.min_pixels, max_pixels=self.max_pixels)
+            all_messages.append(self._build_messages(base64_data, mime_type, prompt))
+
+        results = [None] * len(input_data)
+        for i in tqdm(range(0, len(all_messages), batch_size), desc="Parsing", file=sys.stdout):
+            batch_messages = all_messages[i : i + batch_size]
+            outputs = self._llm.chat(
+                batch_messages,
+                sampling_params=sampling_params,
+                use_tqdm=False,
+                chat_template_kwargs=chat_template_kwargs,
+            )
+            for j, output in enumerate(outputs):
+                results[i + j] = output.outputs[0].text
+
+        return results

infinity_parser2/backends/vllm_server.py

@@ -0,0 +1,148 @@
+"""vLLM Server backend for Infinity-Parser2.
+
+Uses vLLM OpenAI-Compatible Server for online inference.
+"""
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import sys
+from typing import Union
+
+from openai import OpenAI
+from PIL import Image
+from tqdm import tqdm
+
+from .base import BaseBackend
+from ..utils import encode_file_to_base64
+
+
+class VLLMServerBackend(BaseBackend):
+    """Online inference backend using vLLM OpenAI-Compatible Server.
+
+    Sends requests to a running vLLM server via HTTP API.
+    Reference: https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
+    """
+
+    def __init__(
+        self,
+        model_name: str = "infly/Infinity-Parser2-Pro",
+        api_url: str = "http://localhost:8000/v1/chat/completions",
+        api_key: str = "EMPTY",
+        timeout: int = 300,
+        min_pixels: int = 2048,
+        max_pixels: int = 16777216,
+        **kwargs,
+    ):
+        """Initialize vLLM Server backend.
+
+        Args:
+            model_name: Model name (must match server).
+            api_url: Full URL to the chat completions endpoint.
+            api_key: API key for authentication.
+            timeout: Request timeout in seconds.
+            **kwargs: Additional arguments for requests.
+        """
+        device = kwargs.pop("device", "cuda")
+        super().__init__(model_name, device, **kwargs)
+        self.api_url = api_url
+        self.api_key = api_key
+        self.timeout = timeout
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.client = OpenAI(api_key=self.api_key, base_url=self.api_url.rsplit("/chat/completions", 1)[0])
+        self.init()
+
+    def init(self) -> None:
+        """Validate server connection.
+
+        Note: This is a no-op as the server is started separately.
+        Call this to verify connectivity.
+        """
+        try:
+            self.client.chat.completions.create(
+                model=self.model_name,
+                messages=[{"role": "user", "content": "ping"}],
+                max_tokens=1,
+                timeout=5,
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Cannot connect to vLLM server at {self.api_url}. "
+                f"Please ensure the server is running. Error: {e}"
+            )
+
+    def parse_batch(
+        self,
+        input_data: list[Union[str, Image.Image]],
+        prompt: str,
+        batch_size: int = 1,
+        **kwargs,
+    ) -> list[str]:
+        """Parse multiple documents via HTTP API with batched requests.
+
+        Args:
+            input_data: List of file paths or PIL Images.
+            prompt: Prompt text for the model.
+            batch_size: Maximum number of images to process in one batch.
+            **kwargs: Additional arguments.
+
+        Returns:
+            List of parsed text content (one per input in the same order).
+        """
+        if not input_data:
+            return []
+
+        max_tokens = kwargs.get("max_new_tokens", kwargs.get("max_tokens", 32768))
+        temperature = kwargs.get("temperature", 0.0)
+        top_p = kwargs.get("top_p", 1.0)
+        extra_body = {
+            "chat_template_kwargs": {
+                "enable_thinking": False
+            }
+        }
+
+        def parse_one(item: Union[str, Image.Image]) -> str:
+            base64_data, mime_type = encode_file_to_base64(item, min_pixels=self.min_pixels, max_pixels=self.max_pixels)
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_data}"}},
+                        {"type": "text", "text": prompt},
+                    ],
+                }
+            ]
+            response = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                timeout=self.timeout,
+                extra_body=extra_body,
+            )
+            return response.choices[0].message.content
+
+        results: list[str] = [None] * len(input_data)
+        num_batches = (len(input_data) + batch_size - 1) // batch_size
+
+        for batch_idx in range(num_batches):
+            start = batch_idx * batch_size
+            end = min(start + batch_size, len(input_data))
+            batch_items = input_data[start:end]
+
+            with ThreadPoolExecutor(max_workers=len(batch_items)) as executor:
+                future_to_index = {
+                    executor.submit(parse_one, item): start + i
+                    for i, item in enumerate(batch_items)
+                }
+
+                for future in tqdm(as_completed(future_to_index), total=len(batch_items), desc=f"Batch {batch_idx + 1}/{num_batches}", file=sys.stdout):
+                    idx = future_to_index[future]
+                    try:
+                        results[idx] = future.result()
+                    except Exception as e:
+                        raise RuntimeError(
+                            f"Failed to parse input at index {idx}: {e}"
+                        ) from e
+
+        return results