infinity-parser2 0.1.0__py3-none-any.whl

infinity_parser2/cli.py

@@ -0,0 +1,207 @@
+"""Command-line interface for Infinity-Parser2."""
+
+import argparse
+import os
+import sys
+from typing import List, Optional
+
+from . import InfinityParser2
+
+
+def parse_bool(value: str) -> bool:
+    """Convert string to boolean."""
+    if value.lower() in ("true", "1", "yes"):
+        return True
+    elif value.lower() in ("false", "0", "no"):
+        return False
+    raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the argument parser."""
+    parser = argparse.ArgumentParser(
+        prog="parser",
+        description="Infinity-Parser2: Document parsing tool using Infinity-Parser2-Pro model.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Parse a PDF file (default: doc2json -> markdown output)
+  parser document.pdf
+
+  # Parse with doc2md task type
+  parser document.pdf --task doc2md
+
+  # Parse with custom prompt
+  parser document.pdf --task custom --prompt "Extract the title and authors"
+
+  # Parse multiple files
+  parser doc1.pdf doc2.png --output-dir ./results
+
+  # Parse a directory
+  parser ./docs --output-dir ./results
+
+  # Output raw JSON
+  parser document.pdf --output-format json
+
+  # Use transformers backend
+  parser document.pdf --backend transformers
+
+  # Use vllm-server backend
+  parser document.pdf --backend vllm-server --api-url http://localhost:8000/v1/chat/completions
+        """,
+    )
+
+    parser.add_argument(
+        "input",
+        nargs="+",
+        help="Input file(s) or directory path. Supports PDF, PNG, JPG, JPEG, WEBP.",
+    )
+    parser.add_argument(
+        "-o", "--output-dir",
+        default=None,
+        help="Output directory. If not provided, result is printed to stdout.",
+    )
+    parser.add_argument(
+        "--task",
+        default="doc2json",
+        choices=["doc2json", "doc2md", "custom"],
+        help="Parsing task type. Defaults to 'doc2json'.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default=None,
+        help="Custom prompt used only when --task custom.",
+    )
+    parser.add_argument(
+        "--output-format",
+        default="md",
+        choices=["md", "json"],
+        help="Output format. Defaults to 'md'.",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=4,
+        help="Batch size for inference. Defaults to 4.",
+    )
+    parser.add_argument(
+        "--backend",
+        default="vllm-engine",
+        choices=["transformers", "vllm-engine", "vllm-server"],
+        help="Inference backend. Defaults to 'vllm-engine'.",
+    )
+    parser.add_argument(
+        "--model-name",
+        default="infly/Infinity-Parser2-Pro",
+        help="Model name on HuggingFace Hub or local path.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size",
+        type=int,
+        default=None,
+        help="Tensor parallel size for vllm-engine backend.",
+    )
+    parser.add_argument(
+        "--api-url",
+        default="http://localhost:8000/v1/chat/completions",
+        help="API URL for vllm-server backend.",
+    )
+    parser.add_argument(
+        "--api-key",
+        default="EMPTY",
+        help="API key for vllm-server backend.",
+    )
+    parser.add_argument(
+        "--model-cache-dir",
+        default=None,
+        help="Model cache directory.",
+    )
+    parser.add_argument(
+        "--min-pixels",
+        type=int,
+        default=2048,
+        help="Minimum number of pixels for image input (transformers backend only).",
+    )
+    parser.add_argument(
+        "--max-pixels",
+        type=int,
+        default=16777216,
+        help="Maximum number of pixels for image input (transformers backend only).",
+    )
+    parser.add_argument(
+        "--verbose", "-v",
+        action="store_true",
+        help="Print verbose output.",
+    )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version="Infinity-Parser2 0.1.0",
+    )
+
+    return parser
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    """CLI entry point."""
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    input_paths = args.input
+    if len(input_paths) == 1 and os.path.isdir(input_paths[0]):
+        input_data = input_paths[0]
+    else:
+        input_data = input_paths
+
+    if args.verbose:
+        print(f"[Infinity-Parser2] Backend: {args.backend}")
+        print(f"[Infinity-Parser2] Model: {args.model_name}")
+        print(f"[Infinity-Parser2] Task: {args.task}")
+        print(f"[Infinity-Parser2] Input: {input_data}")
+
+    try:
+        parser_client = InfinityParser2(
+            model_name=args.model_name,
+            backend=args.backend,
+            tensor_parallel_size=args.tensor_parallel_size,
+            api_url=args.api_url,
+            api_key=args.api_key,
+            min_pixels=args.min_pixels,
+            max_pixels=args.max_pixels,
+            model_cache_dir=args.model_cache_dir,
+        )
+
+        result = parser_client.parse(
+            input_data=input_data,
+            task_type=args.task,
+            custom_prompt=args.prompt,
+            batch_size=args.batch_size,
+            output_dir=args.output_dir,
+            output_format=args.output_format,
+        )
+
+        if result is not None:
+            if isinstance(result, dict):
+                for path, content in result.items():
+                    print(f"=== {path} ===")
+                    print(content)
+            elif isinstance(result, list):
+                for item in result:
+                    print(item)
+            else:
+                print(result)
+        elif args.verbose:
+            print("[Infinity-Parser2] Results saved to output directory.")
+
+        return 0
+
+    except Exception as e:
+        print(f"[Infinity-Parser2] Error: {e}", file=sys.stderr)
+        if args.verbose:
+            import traceback
+            traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

infinity_parser2/parser.py

@@ -0,0 +1,278 @@
+"""Infinity-Parser2 main interface."""
+
+import os
+import re
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import torch
+from PIL import Image
+
+from .backends import (
+    BaseBackend,
+    TransformersBackend,
+    VLLMEngineBackend,
+    VLLMServerBackend,
+)
+from .prompts import PROMPT_DOC2JSON, PROMPT_DOC2MD, SUPPORTED_TASK_TYPES
+from .utils import *
+
+
+BACKEND_REGISTRY = {
+    "transformers": TransformersBackend,
+    "vllm-engine": VLLMEngineBackend,
+    "vllm-server": VLLMServerBackend,
+}
+
+
+class InfinityParser2:
+    """Document parser using Infinity-Parser2-Pro model.
+
+    Supports parsing of PDF files and images (PNG, JPG, etc.) into structured text.
+
+    Args:
+        model_name: Model name on HuggingFace Hub (e.g., "infly/Infinity-Parser2-Pro")
+            or local path to a downloaded model. Defaults to "infly/Infinity-Parser2-Pro".
+        backend: Inference backend. Options:
+            - "transformers": HuggingFace transformers (local inference)
+            - "vllm-engine": vLLM Engine (local batch inference via LLM class)
+            - "vllm-server": vLLM OpenAI-Compatible Server (HTTP API)
+          Defaults to "vllm-engine".
+        tensor_parallel_size: Tensor parallel size for vllm-engine.
+            Defaults to the number of available GPUs (via torch.cuda.device_count()).
+        device: Device type, must be "cuda". Raises ValueError if set to anything else.
+        api_url: API URL for vllm-server backend.
+        api_key: API key for vllm-server backend.
+        min_pixels: Minimum number of pixels for image input (transformers backend only).
+            Defaults to 2048.
+        max_pixels: Maximum number of pixels for image input (transformers backend only).
+            Defaults to 16777216 (~4096x4096).
+        **kwargs: Additional arguments passed to the backend.
+
+    Example:
+        >>> from infinity_parser2 import InfinityParser2
+        >>> parser = InfinityParser2(model_name="infly/Infinity-Parser2-Pro")
+        >>> result = parser.parse("document.pdf")
+    """
+
+    def __init__(
+        self,
+        model_name: str = "infly/Infinity-Parser2-Pro",
+        backend: str = "vllm-engine",
+        tensor_parallel_size: Optional[int] = None,
+        device: str = "cuda",
+        api_url: str = "http://localhost:8000/v1/chat/completions",
+        api_key: str = "EMPTY",
+        min_pixels: int = 2048,
+        max_pixels: int = 16777216,
+        model_cache_dir: Optional[str] = None,
+        **kwargs,
+    ):
+        if device != "cuda":
+            raise ValueError("device must be 'cuda' for Infinity-Parser2-Pro.")
+
+        self.model_name = model_name
+        self.backend_name = backend.lower()
+        self.tensor_parallel_size = (
+            tensor_parallel_size
+            if tensor_parallel_size is not None
+            else torch.cuda.device_count()
+        )
+        self.device = device
+        self.api_url = api_url
+        self.api_key = api_key
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.kwargs = kwargs
+
+        # Initialize model cache and resolve model path (stored separately)
+        cache = get_model_cache(model_cache_dir)
+        self._model_path = cache.resolve_model_path(self.model_name)
+
+        self._backend: BaseBackend = self._init_backend()
+
+    def _init_backend(self) -> BaseBackend:
+        """Initialize and return the backend instance."""
+        if self.backend_name not in BACKEND_REGISTRY:
+            raise ValueError(
+                f"Unsupported backend: {self.backend_name}. "
+                f"Supported backends: {list(BACKEND_REGISTRY.keys())}"
+            )
+        backend_cls = BACKEND_REGISTRY[self.backend_name]
+        common_kwargs = {
+            "model_name": self._model_path,
+            "device": self.device,
+            "min_pixels": self.min_pixels,
+            "max_pixels": self.max_pixels,
+            **self.kwargs,
+        }
+        if self.backend_name == "vllm-server":
+            backend_kwargs = {**common_kwargs, "api_url": self.api_url, "api_key": self.api_key}
+        elif self.backend_name == "vllm-engine":
+            backend_kwargs = {**common_kwargs, "tensor_parallel_size": self.tensor_parallel_size}
+        else:  # transformers
+            backend_kwargs = common_kwargs
+        return backend_cls(**backend_kwargs)
+
+    def _resolve_prompt(self, task_type: str, custom_prompt: Optional[str]) -> str:
+        """Resolve the prompt to use based on task_type and custom_prompt.
+
+        Args:
+            task_type: The task type (e.g., "doc2json", "doc2md", "custom").
+            custom_prompt: Custom prompt, only used when task_type is "custom".
+
+        Returns:
+            The resolved prompt string.
+        """
+        if task_type == "custom":
+            assert custom_prompt is not None, "custom_prompt must be provided when task_type='custom'"
+            return custom_prompt
+        if task_type == "doc2json":
+            return PROMPT_DOC2JSON
+        if task_type == "doc2md":
+            return PROMPT_DOC2MD
+        # Fallback for unknown task types (should not happen with proper validation)
+        return "Please transform the document's contents into Markdown format."
+
+    def parse(
+        self,
+        input_data: Union[str, List[str], Image.Image],
+        task_type: str = "doc2json",
+        custom_prompt: Optional[str] = None,
+        batch_size: int = 4,
+        output_dir: Optional[str] = None,
+        output_format: str = "md",
+        **kwargs,
+    ) -> Optional[Union[str, List[str], Dict[str, str]]]:
+        """Parse document(s) and extract text content.
+
+        Args:
+            input_data: Input can be:
+                - str: Single file path or directory path
+                - List[str]: List of file paths
+                - PIL.Image.Image: Image object
+            task_type: Parsing task type. Options:
+                - "doc2json": Extract layout to JSON, return JSON string.
+                - "doc2md": Directly convert to Markdown, return Markdown.
+                - "custom": Use custom_prompt for parsing.
+                Defaults to "doc2json".
+            custom_prompt: Custom prompt text for the model. Used only when
+                task_type is "custom". Defaults to None.
+            batch_size: Number of images to process in one batch. Defaults to 4.
+            output_dir: If provided, results are saved to output_dir and this function
+                returns None. If None, results are returned directly.
+            output_format: Output format for results. Options: "md" or "json".
+                Defaults to "md".
+                - For doc2json tasks:
+                    - output_format="md": Returns markdown (converts JSON to markdown
+                      via convert_json_to_markdown). If output_dir is set, saves only
+                      the markdown result.
+                    - output_format="json": Returns raw JSON result. If output_dir is
+                      set, saves only the JSON result.
+                - For doc2md tasks or custom prompts: Only "md" is supported.
+                  If "json" is passed, a ValueError will be raised.
+            **kwargs: Additional arguments passed to the model.
+
+        Returns:
+            When output_dir is None:
+                - str: Parsed result for a single file or image.
+                - List[str]: Parsed results for a list of files.
+                - Dict[str, str]: Mapping from file path to parsed result for a directory.
+            When output_dir is set, returns None.
+
+        Example:
+            >>> parser = InfinityParser2()
+            >>> # Single file, returns str
+            >>> result = parser.parse("document.pdf")
+            >>> # Multiple files, returns List[str]
+            >>> result = parser.parse(["doc1.pdf", "doc2.pdf"])
+            >>> # Directory, returns Dict[str, str]
+            >>> result = parser.parse("/path/to/docs")
+            >>> # Save results to output_dir, returns None
+            >>> parser.parse("document.pdf", output_dir="./output")
+        """
+        if task_type not in SUPPORTED_TASK_TYPES:
+            raise ValueError(f"task_type must be one of {SUPPORTED_TASK_TYPES}, got '{task_type}'")
+
+        if output_format not in SUPPORTED_OUTPUT_FORMATS:
+            raise ValueError(f"output_format must be one of {SUPPORTED_OUTPUT_FORMATS}, got '{output_format}'")
+
+        if output_format == "json" and task_type != "doc2json":
+            raise ValueError(
+                "output_format='json' is only supported for doc2json tasks. "
+                "For other task types, output_format must be 'md'."
+            )
+
+        prompt = self._resolve_prompt(task_type, custom_prompt)
+
+        is_directory = isinstance(input_data, str) and os.path.isdir(input_data)
+        file_paths = normalize_input(input_data)
+        file_results = self._parse_files(
+            file_paths, prompt, task_type, batch_size, output_format, **kwargs
+        )
+
+        if output_dir is not None:
+            save_results(
+                file_paths, file_results, output_dir,
+                task_type=task_type, output_format=output_format
+            )
+        elif is_directory:
+            return dict(zip(file_paths, file_results))
+        elif len(file_results) == 1:
+            return file_results[0]
+        else:
+            return file_results
+
+    def _parse_files(
+        self,
+        inputs: List[Union[str, Image.Image]],
+        prompt: Optional[str],
+        task_type: str,
+        batch_size: int = 4,
+        output_format: str = "md",
+        **kwargs,
+    ) -> List[str]:
+        """Parse multiple files with batched inference.
+
+        All images (including PDF pages) are collected and batched together for
+        efficient inference. Results are then aggregated back to the original
+        file-level granularity.
+        """
+
+        # prepare batch entries
+        batch_entries = prepare_batch_entries(inputs)
+        if not batch_entries:
+            return [] if len(inputs) > 1 else ""
+
+        # parse batch
+        raw_inputs = [entry[1] for entry in batch_entries]
+        batch_results = self._backend.parse_batch(raw_inputs, prompt, batch_size=batch_size, **kwargs)
+
+        # aggregate batch results
+        num_files = len({entry[0] for entry in batch_entries})
+        page_results: List[List[str]] = [[] for _ in range(num_files)]
+        file_results: List[str] = [""] * num_files
+
+        for entry_idx, (file_idx, image_input) in enumerate(batch_entries):
+            raw_result = batch_results[entry_idx]
+
+            # postprocess result
+            if task_type == "doc2json":
+                text = postprocess_doc2json_result(raw_result, image_input, output_format)
+            elif task_type == "doc2md":
+                text = postprocess_doc2md_result(raw_result)
+            else:
+                text = raw_result
+
+            page_results[file_idx].append(text)
+
+        # Join results based on length of page_results and output_format
+        for idx in range(num_files):
+            if len(page_results[idx]) == 1:
+                file_results[idx] = page_results[idx][0]
+            elif output_format == "json":
+                file_results[idx] = "[" + ",".join(page_results[idx]) + "]"
+            else:
+                file_results[idx] = "\n\n".join(page_results[idx])
+
+        return file_results

infinity_parser2/prompts.py

@@ -0,0 +1,57 @@
+"""Prompts for Infinity-Parser2."""
+
+__all__ = [
+    "PROMPT_DOC2JSON",
+    "PROMPT_DOC2MD",
+    "SUPPORTED_TASK_TYPES",
+]
+
+SUPPORTED_TASK_TYPES = ["doc2json", "doc2md", "custom"]
+
+
+# doc2json prompt (outputs JSON format)
+PROMPT_DOC2JSON = """
+- Extract layout information from the provided PDF image.
+- For each layout element, output its bbox, category, and the text content within the bbox.
+- Bbox format: [x1, y1, x2, y2].
+- Allowed layout categories: ['header', 'title', 'text', 'figure', 'table', 'formula',
+  'figure_caption', 'table_caption', 'formula_caption', 'figure_footnote',
+  'table_footnote', 'page_footnote', 'footer'].
+- Text extraction and formatting:
+  1) For 'figure', the text field must be an empty string.
+  2) For 'formula', format text as LaTeX.
+  3) For 'table', format text as HTML.
+  4) For all other categories (e.g., text, title), format text as Markdown.
+- The output text must be exactly the original text from the image,
+  with no translation or rewriting.
+- Sort all layout elements in human reading order.
+- Final output must be a single JSON object.
+"""
+
+# doc2md prompt (outputs Markdown format directly)
+PROMPT_DOC2MD = """
+You are an AI assistant specialized in converting PDF images to Markdown format. Please follow these instructions for the conversion:
+
+1. Text Processing:
+- Accurately recognize all text content in the PDF image without guessing or inferring.
+- Convert the recognized text into Markdown format.
+- Maintain the original document structure, including headings, paragraphs, lists, etc.
+
+2. Mathematical Formula Processing:
+- Convert all mathematical formulas to LaTeX format.
+- Enclose inline formulas with $ $. For example: This is an inline formula $E = mc^2$
+- Enclose block formulas with $$ $$. For example: $$\\frac{-b \\pm \\sqrt{b^2 - 4ac}}{2a}$$
+
+3. Table Processing:
+- Convert tables to HTML format.
+
+4. Figure Handling:
+- Ignore figures content in the PDF image. Do not attempt to describe or convert images.
+
+5. Output Format:
+- Ensure the output Markdown document has a clear structure with appropriate line breaks between elements.
+- For complex layouts, try to maintain the original document's structure and format as closely as possible.
+
+Please strictly follow these guidelines to ensure accuracy and consistency in the conversion. Your task is to accurately convert the content of the PDF image into Markdown format without adding any extra explanations or comments.
+"""
+

infinity_parser2/utils/__init__.py

@@ -0,0 +1,43 @@
+"""Utility functions for Infinity-Parser2."""
+
+from .file import (
+    get_files_from_directory,
+    is_supported_file,
+    normalize_input,
+    prepare_batch_entries,
+    save_results,
+    SUPPORTED_OUTPUT_FORMATS,
+)
+from .image import encode_file_to_base64, load_image
+from .model import ModelCache, get_model_cache
+from .pdf import convert_pdf_to_images
+from .utils import (
+    convert_json_to_markdown,
+    extract_json_content,
+    obtain_origin_hw,
+    postprocess_doc2json_result,
+    restore_abs_bbox_coordinates,
+    postprocess_doc2md_result,
+    truncate_last_incomplete_element,
+)
+
+__all__ = [
+    "convert_pdf_to_images",
+    "convert_json_to_markdown",
+    "extract_json_content",
+    "encode_file_to_base64",
+    "get_files_from_directory",
+    "get_model_cache",
+    "is_supported_file",
+    "load_image",
+    "ModelCache",
+    "normalize_input",
+    "obtain_origin_hw",
+    "postprocess_doc2json_result",
+    "postprocess_doc2md_result",
+    "prepare_batch_entries",
+    "restore_abs_bbox_coordinates",
+    "save_results",
+    "SUPPORTED_OUTPUT_FORMATS",
+    "truncate_last_incomplete_element",
+]