In [ ]

import json
import time
from pathlib import Path

import json import time from pathlib import Path

In [ ]

from docling_core.types.doc import DocItemLabel, ImageRefMode
from docling_core.types.doc.document import DEFAULT_EXPORT_LABELS

from docling_core.types.doc import DocItemLabel, ImageRefMode from docling_core.types.doc.document import DEFAULT_EXPORT_LABELS

In [ ]

from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import (
    VlmPipelineOptions,
    smoldocling_vlm_mlx_conversion_options,
)
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.pipeline.vlm_pipeline import VlmPipeline

from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import ( VlmPipelineOptions, smoldocling_vlm_mlx_conversion_options, ) from docling.document_converter import DocumentConverter, PdfFormatOption from docling.pipeline.vlm_pipeline import VlmPipeline

In [ ]

sources = [
    # "tests/data/2305.03393v1-pg9-img.png",
    "tests/data/pdf/2305.03393v1-pg9.pdf",
]

sources = [ # "tests/data/2305.03393v1-pg9-img.png", "tests/data/pdf/2305.03393v1-pg9.pdf", ]

In [ ]

## Use experimental VlmPipeline
pipeline_options = VlmPipelineOptions()
# If force_backend_text = True, text from backend will be used instead of generated text
pipeline_options.force_backend_text = False

## 使用 experimental VlmPipeline pipeline_options = VlmPipelineOptions() # 如果 force_backend_text = True，将使用来自后端的文本而不是生成的文本 pipeline_options.force_backend_text = False

在 GPU 系统上，通过 CUDA 启用 flash_attention_2：¶

pipeline_options.accelerator_options.device = AcceleratorDevice.CUDA pipeline_options.accelerator_options.cuda_use_flash_attention2 = True

选择一个 VLM 模型。默认选择 SmolDocling-256M¶

pipeline_options.vlm_options = smoldocling_vlm_conversion_options

In [ ]

## Pick a VLM model. Fast Apple Silicon friendly implementation for SmolDocling-256M via MLX
pipeline_options.vlm_options = smoldocling_vlm_mlx_conversion_options

## 选择一个 VLM 模型。通过 MLX 实现的适用于 Apple Silicon 的 SmolDocling-256M 快速实现 pipeline_options.vlm_options = smoldocling_vlm_mlx_conversion_options

备选 VLM 模型：¶

pipeline_options.vlm_options = granite_vision_vlm_conversion_options

In [ ]

## Set up pipeline for PDF or image inputs
converter = DocumentConverter(
    format_options={
        InputFormat.PDF: PdfFormatOption(
            pipeline_cls=VlmPipeline,
            pipeline_options=pipeline_options,
        ),
        InputFormat.IMAGE: PdfFormatOption(
            pipeline_cls=VlmPipeline,
            pipeline_options=pipeline_options,
        ),
    }
)

## 为 PDF 或图像输入设置流水线 converter = DocumentConverter( format_options={ InputFormat.PDF: PdfFormatOption( pipeline_cls=VlmPipeline, pipeline_options=pipeline_options, ), InputFormat.IMAGE: PdfFormatOption( pipeline_cls=VlmPipeline, pipeline_options=pipeline_options, ), } )

In [ ]

out_path = Path("scratch")
out_path.mkdir(parents=True, exist_ok=True)

out_path = Path("scratch") out_path.mkdir(parents=True, exist_ok=True)

In [ ]

for source in sources:
    start_time = time.time()
    print("================================================")
    print(f"Processing... {source}")
    print("================================================")
    print("")

    res = converter.convert(source)

    print("")
    print(res.document.export_to_markdown())

    for page in res.pages:
        print("")
        print("Predicted page in DOCTAGS:")
        print(page.predictions.vlm_response.text)

    res.document.save_as_html(
        filename=Path(f"{out_path}/{res.input.file.stem}.html"),
        image_mode=ImageRefMode.REFERENCED,
        labels=[*DEFAULT_EXPORT_LABELS, DocItemLabel.FOOTNOTE],
    )

    with (out_path / f"{res.input.file.stem}.json").open("w") as fp:
        fp.write(json.dumps(res.document.export_to_dict()))

    res.document.save_as_json(
        out_path / f"{res.input.file.stem}.json",
        image_mode=ImageRefMode.PLACEHOLDER,
    )

    res.document.save_as_markdown(
        out_path / f"{res.input.file.stem}.md",
        image_mode=ImageRefMode.PLACEHOLDER,
    )

    pg_num = res.document.num_pages()
    print("")
    inference_time = time.time() - start_time
    print(
        f"Total document prediction time: {inference_time:.2f} seconds, pages: {pg_num}"
    )

for source in sources: start_time = time.time() print("================================================") print(f"正在处理... {source}") print("================================================") print("") res = converter.convert(source) print("") print(res.document.export_to_markdown()) for page in res.pages: print("") print("DOCTAGS 中的预测页面：") print(page.predictions.vlm_response.text) res.document.save_as_html( filename=Path(f"{out_path}/{res.input.file.stem}.html"), image_mode=ImageRefMode.REFERENCED, labels=[*DEFAULT_EXPORT_LABELS, DocItemLabel.FOOTNOTE], ) with (out_path / f"{res.input.file.stem}.json").open("w") as fp: fp.write(json.dumps(res.document.export_to_dict())) res.document.save_as_json( out_path / f"{res.input.file.stem}.json", image_mode=ImageRefMode.PLACEHOLDER, ) res.document.save_as_markdown( out_path / f"{res.input.file.stem}.md", image_mode=ImageRefMode.PLACEHOLDER, ) pg_num = res.document.num_pages() print("") inference_time = time.time() - start_time print( f"总文档预测时间：{inference_time:.2f} 秒，页数：{pg_num}" )

In [ ]

print("================================================")
print("done!")
print("================================================")

print("================================================") print("完成！") print("================================================")