Demo layout vlm

In [ ]:

"""Demo script for the new ThreadedLayoutVlmPipeline.

This script demonstrates the usage of the experimental ThreadedLayoutVlmPipeline pipeline
that combines layout model preprocessing with VLM processing in a threaded manner.
"""

"""Demo script for the new ThreadedLayoutVlmPipeline. This script demonstrates the usage of the experimental ThreadedLayoutVlmPipeline pipeline that combines layout model preprocessing with VLM processing in a threaded manner. """

In [ ]:

import argparse
import logging
import traceback
from pathlib import Path

import argparse import logging import traceback from pathlib import Path

In [ ]:

from docling.datamodel.base_models import ConversionStatus, InputFormat
from docling.datamodel.pipeline_options_vlm_model import ApiVlmOptions, ResponseFormat
from docling.datamodel.vlm_model_specs import GRANITEDOCLING_TRANSFORMERS
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.experimental.datamodel.threaded_layout_vlm_pipeline_options import (
    ThreadedLayoutVlmPipelineOptions,
)
from docling.experimental.pipeline.threaded_layout_vlm_pipeline import (
    ThreadedLayoutVlmPipeline,
)

from docling.datamodel.base_models import ConversionStatus, InputFormat from docling.datamodel.pipeline_options_vlm_model import ApiVlmOptions, ResponseFormat from docling.datamodel.vlm_model_specs import GRANITEDOCLING_TRANSFORMERS from docling.document_converter import DocumentConverter, PdfFormatOption from docling.experimental.datamodel.threaded_layout_vlm_pipeline_options import ( ThreadedLayoutVlmPipelineOptions, ) from docling.experimental.pipeline.threaded_layout_vlm_pipeline import ( ThreadedLayoutVlmPipeline, )

In [ ]:

_log = logging.getLogger(__name__)

_log = logging.getLogger(__name__)

In [ ]:

def _parse_args():
    parser = argparse.ArgumentParser(
        description="Demo script for the experimental ThreadedLayoutVlmPipeline"
    )
    parser.add_argument(
        "--input-file",
        type=str,
        default="tests/data/pdf/code_and_formula.pdf",
        help="Path to a PDF file",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="scratch/demo_layout_vlm/",
        help="Output directory for converted files",
    )
    return parser.parse_args()

def _parse_args(): parser = argparse.ArgumentParser( description="Demo script for the experimental ThreadedLayoutVlmPipeline" ) parser.add_argument( "--input-file", type=str, default="tests/data/pdf/code_and_formula.pdf", help="Path to a PDF file", ) parser.add_argument( "--output-dir", type=str, default="scratch/demo_layout_vlm/", help="Output directory for converted files", ) return parser.parse_args()

Can be used to read multiple pdf files under a folder def _get_docs(input_doc_path): """Yield DocumentStream objects from list of input document paths""" for path in input_doc_path: buf = BytesIO(path.read_bytes()) stream = DocumentStream(name=path.name, stream=buf) yield stream

In [ ]:

def openai_compatible_vlm_options(
    model: str,
    prompt: str,
    format: ResponseFormat,
    hostname_and_port,
    temperature: float = 0.7,
    max_tokens: int = 4096,
    api_key: str = "",
    skip_special_tokens=False,
):
    headers = {}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    options = ApiVlmOptions(
        url=f"http://{hostname_and_port}/v1/chat/completions",  # LM studio defaults to port 1234, VLLM to 8000
        params=dict(
            model=model,
            max_tokens=max_tokens,
            skip_special_tokens=skip_special_tokens,  # needed for VLLM
        ),
        headers=headers,
        prompt=prompt,
        timeout=90,
        scale=2.0,
        temperature=temperature,
        response_format=format,
    )

    return options

def openai_compatible_vlm_options( model: str, prompt: str, format: ResponseFormat, hostname_and_port, temperature: float = 0.7, max_tokens: int = 4096, api_key: str = "", skip_special_tokens=False, ): headers = {} if api_key: headers["Authorization"] = f"Bearer {api_key}" options = ApiVlmOptions( url=f"http://{hostname_and_port}/v1/chat/completions", # LM studio defaults to port 1234, VLLM to 8000 params=dict( model=model, max_tokens=max_tokens, skip_special_tokens=skip_special_tokens, # needed for VLLM ), headers=headers, prompt=prompt, timeout=90, scale=2.0, temperature=temperature, response_format=format, ) return options

In [ ]:

def demo_threaded_layout_vlm_pipeline(
    input_doc_path: Path, out_dir_layout_aware: Path, use_api_vlm: bool
):
    """Demonstrate the threaded layout+VLM pipeline."""

    vlm_options = GRANITEDOCLING_TRANSFORMERS.model_copy()

    if use_api_vlm:
        vlm_options = openai_compatible_vlm_options(
            model="granite-docling-258m-mlx",  # For VLLM use "ibm-granite/granite-docling-258M"
            hostname_and_port="localhost:1234",  # LM studio defaults to port 1234, VLLM to 8000
            prompt="Convert this page to docling.",
            format=ResponseFormat.DOCTAGS,
            api_key="",
        )
    vlm_options.track_input_prompt = True

    # Configure pipeline options
    print("Configuring pipeline options...")
    pipeline_options_layout_aware = ThreadedLayoutVlmPipelineOptions(
        # VLM configuration - defaults to GRANITEDOCLING_TRANSFORMERS
        vlm_options=vlm_options,
        # Layout configuration - defaults to DOCLING_LAYOUT_HERON
        # Batch sizes for parallel processing
        layout_batch_size=2,
        vlm_batch_size=1,
        # Queue configuration
        queue_max_size=10,
        # Image processing
        images_scale=vlm_options.scale,
        generate_page_images=True,
        enable_remote_services=use_api_vlm,
    )

    # Create converter with the new pipeline
    print("Initializing DocumentConverter (this may take a while - loading models)...")
    doc_converter_layout_enhanced = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(
                pipeline_cls=ThreadedLayoutVlmPipeline,
                pipeline_options=pipeline_options_layout_aware,
            )
        }
    )

    result_layout_aware = doc_converter_layout_enhanced.convert(
        source=input_doc_path, raises_on_error=False
    )

    if result_layout_aware.status == ConversionStatus.FAILURE:
        _log.error(f"Conversion failed: {result_layout_aware.status}")

    doc_filename = result_layout_aware.input.file.stem
    result_layout_aware.document.save_as_json(
        out_dir_layout_aware / f"{doc_filename}.json"
    )

    result_layout_aware.document.save_as_html(
        out_dir_layout_aware / f"{doc_filename}.html", split_page_view=True
    )
    for page in result_layout_aware.pages:
        _log.info("Page %s of VLM response:", page.page_no)
        if page.predictions.vlm_response:
            _log.info(page.predictions.vlm_response)

def demo_threaded_layout_vlm_pipeline( input_doc_path: Path, out_dir_layout_aware: Path, use_api_vlm: bool ): """Demonstrate the threaded layout+VLM pipeline.""" vlm_options = GRANITEDOCLING_TRANSFORMERS.model_copy() if use_api_vlm: vlm_options = openai_compatible_vlm_options( model="granite-docling-258m-mlx", # For VLLM use "ibm-granite/granite-docling-258M" hostname_and_port="localhost:1234", # LM studio defaults to port 1234, VLLM to 8000 prompt="Convert this page to docling.", format=ResponseFormat.DOCTAGS, api_key="", ) vlm_options.track_input_prompt = True # Configure pipeline options print("Configuring pipeline options...") pipeline_options_layout_aware = ThreadedLayoutVlmPipelineOptions( # VLM configuration - defaults to GRANITEDOCLING_TRANSFORMERS vlm_options=vlm_options, # Layout configuration - defaults to DOCLING_LAYOUT_HERON # Batch sizes for parallel processing layout_batch_size=2, vlm_batch_size=1, # Queue configuration queue_max_size=10, # Image processing images_scale=vlm_options.scale, generate_page_images=True, enable_remote_services=use_api_vlm, ) # Create converter with the new pipeline print("Initializing DocumentConverter (this may take a while - loading models)...") doc_converter_layout_enhanced = DocumentConverter( format_options={ InputFormat.PDF: PdfFormatOption( pipeline_cls=ThreadedLayoutVlmPipeline, pipeline_options=pipeline_options_layout_aware, ) } ) result_layout_aware = doc_converter_layout_enhanced.convert( source=input_doc_path, raises_on_error=False ) if result_layout_aware.status == ConversionStatus.FAILURE: _log.error(f"Conversion failed: {result_layout_aware.status}") doc_filename = result_layout_aware.input.file.stem result_layout_aware.document.save_as_json( out_dir_layout_aware / f"{doc_filename}.json" ) result_layout_aware.document.save_as_html( out_dir_layout_aware / f"{doc_filename}.html", split_page_view=True ) for page in result_layout_aware.pages: _log.info("Page %s of VLM response:", page.page_no) if page.predictions.vlm_response: _log.info(page.predictions.vlm_response)

In [ ]:

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    try:
        args = _parse_args()
        _log.info(
            f"Parsed arguments: input={args.input_file}, output={args.output_dir}"
        )

        input_path = Path(args.input_file)

        if not input_path.exists():
            raise FileNotFoundError(f"Input file does not exist: {input_path}")

        if input_path.suffix.lower() != ".pdf":
            raise ValueError(f"Input file must be a PDF: {input_path}")

        out_dir_layout_aware = Path(args.output_dir) / "layout_aware/"
        out_dir_layout_aware.mkdir(parents=True, exist_ok=True)

        use_api_vlm = False  # Set to False to use inline VLM model

        demo_threaded_layout_vlm_pipeline(input_path, out_dir_layout_aware, use_api_vlm)
    except Exception:
        traceback.print_exc()
        raise

if __name__ == "__main__": logging.basicConfig(level=logging.INFO) try: args = _parse_args() _log.info( f"Parsed arguments: input={args.input_file}, output={args.output_dir}" ) input_path = Path(args.input_file) if not input_path.exists(): raise FileNotFoundError(f"Input file does not exist: {input_path}") if input_path.suffix.lower() != ".pdf": raise ValueError(f"Input file must be a PDF: {input_path}") out_dir_layout_aware = Path(args.output_dir) / "layout_aware/" out_dir_layout_aware.mkdir(parents=True, exist_ok=True) use_api_vlm = False # Set to False to use inline VLM model demo_threaded_layout_vlm_pipeline(input_path, out_dir_layout_aware, use_api_vlm) except Exception: traceback.print_exc() raise