mirror of
https://github.com/DS4SD/docling.git
synced 2025-07-26 20:14:47 +00:00
222 lines
8.0 KiB
Python
222 lines
8.0 KiB
Python
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from docling_core.types.doc import DocItemLabel, ImageRefMode
|
|
from docling_core.types.doc.document import DEFAULT_EXPORT_LABELS
|
|
|
|
from docling.datamodel.base_models import InputFormat
|
|
from docling.datamodel.pipeline_model_specializations import (
|
|
HuggingFaceVlmOptions,
|
|
InferenceFramework,
|
|
ResponseFormat,
|
|
granite_vision_vlm_conversion_options,
|
|
granite_vision_vlm_ollama_conversion_options,
|
|
phi_vlm_conversion_options,
|
|
pixtral_12b_vlm_conversion_options,
|
|
pixtral_12b_vlm_mlx_conversion_options,
|
|
qwen25_vl_3b_vlm_mlx_conversion_options,
|
|
smoldocling_vlm_conversion_options,
|
|
smoldocling_vlm_mlx_conversion_options,
|
|
)
|
|
from docling.datamodel.pipeline_options import (
|
|
VlmPipelineOptions,
|
|
)
|
|
from docling.document_converter import DocumentConverter, PdfFormatOption
|
|
from docling.pipeline.vlm_pipeline import VlmPipeline
|
|
|
|
from tabulate import tabulate
|
|
|
|
## 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
|
|
pipeline_options.generate_page_images = True
|
|
|
|
## On GPU systems, enable flash_attention_2 with CUDA:
|
|
# pipeline_options.accelerator_options.device = AcceleratorDevice.CUDA
|
|
# pipeline_options.accelerator_options.cuda_use_flash_attention2 = True
|
|
|
|
## Pick a VLM model. We choose SmolDocling-256M by default
|
|
# pipeline_options.vlm_options = smoldocling_vlm_conversion_options
|
|
|
|
## Pick a VLM model. Fast Apple Silicon friendly implementation for SmolDocling-256M via MLX
|
|
# pipeline_options.vlm_options = smoldocling_vlm_mlx_conversion_options
|
|
|
|
## Alternative VLM models:
|
|
# pipeline_options.vlm_options = granite_vision_vlm_conversion_options
|
|
|
|
pipeline_options.vlm_options = phi_vlm_conversion_options
|
|
# pipeline_options.vlm_options = qwen25_vl_3b_vlm_mlx_conversion_options
|
|
|
|
"""
|
|
pixtral_vlm_conversion_options = HuggingFaceVlmOptions(
|
|
repo_id="mistralai/Pixtral-12B-Base-2409",
|
|
prompt="OCR this image and export it in MarkDown.",
|
|
response_format=ResponseFormat.MARKDOWN,
|
|
inference_framework=InferenceFramework.TRANSFORMERS_LlavaForConditionalGeneration,
|
|
)
|
|
pipeline_options.vlm_options = pixtral_vlm_conversion_options
|
|
"""
|
|
|
|
"""
|
|
pixtral_vlm_conversion_options = HuggingFaceVlmOptions(
|
|
repo_id="mistral-community/pixtral-12b",
|
|
prompt="OCR this image and export it in MarkDown.",
|
|
response_format=ResponseFormat.MARKDOWN,
|
|
inference_framework=InferenceFramework.TRANSFORMERS_LlavaForConditionalGeneration,
|
|
)
|
|
pipeline_options.vlm_options = pixtral_vlm_conversion_options
|
|
"""
|
|
|
|
"""
|
|
phi_vlm_conversion_options = HuggingFaceVlmOptions(
|
|
repo_id="microsoft/Phi-4-multimodal-instruct",
|
|
# prompt="OCR the full page to markdown.",
|
|
prompt="OCR this image and export it in MarkDown.",
|
|
response_format=ResponseFormat.MARKDOWN,
|
|
inference_framework=InferenceFramework.TRANSFORMERS_AutoModelForCausalLM,
|
|
)
|
|
pipeline_options.vlm_options = phi_vlm_conversion_options
|
|
"""
|
|
|
|
"""
|
|
pixtral_vlm_conversion_options = HuggingFaceVlmOptions(
|
|
repo_id="mlx-community/pixtral-12b-bf16",
|
|
prompt="Convert this page to markdown. Do not miss any text and only output the bare MarkDown!",
|
|
response_format=ResponseFormat.MARKDOWN,
|
|
inference_framework=InferenceFramework.MLX,
|
|
scale=1.0,
|
|
)
|
|
pipeline_options.vlm_options = pixtral_vlm_conversion_options
|
|
"""
|
|
|
|
"""
|
|
qwen_vlm_conversion_options = HuggingFaceVlmOptions(
|
|
repo_id="mlx-community/Qwen2.5-VL-3B-Instruct-bf16",
|
|
prompt="Convert this full page to markdown. Do not miss any text and only output the bare MarkDown!",
|
|
response_format=ResponseFormat.MARKDOWN,
|
|
inference_framework=InferenceFramework.MLX,
|
|
)
|
|
pipeline_options.vlm_options = qwen_vlm_conversion_options
|
|
"""
|
|
|
|
def convert(sources: list[Path], converter):
|
|
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())
|
|
|
|
model_id = pipeline_options.vlm_options.repo_id.replace("/", "_")
|
|
framework = pipeline_options.vlm_options.inference_framework
|
|
fname = f"{res.input.file.stem}-{model_id}-{framework}"
|
|
|
|
inference_time = 0.0
|
|
for i, page in enumerate(res.pages):
|
|
inference_time += page.predictions.vlm_response.generation_time
|
|
print("")
|
|
print(
|
|
f" ---------- Predicted page {i} in {pipeline_options.vlm_options.response_format} in {page.predictions.vlm_response.generation_time} [sec]:"
|
|
)
|
|
print(page.predictions.vlm_response.text)
|
|
print(" ---------- ")
|
|
|
|
print("===== Final output of the converted document =======")
|
|
|
|
with (out_path / f"{fname}.json").open("w") as fp:
|
|
fp.write(json.dumps(res.document.export_to_dict()))
|
|
|
|
res.document.save_as_json(
|
|
out_path / f"{fname}.json",
|
|
image_mode=ImageRefMode.PLACEHOLDER,
|
|
)
|
|
print(f" => produced {out_path / fname}.json")
|
|
|
|
res.document.save_as_markdown(
|
|
out_path / f"{fname}.md",
|
|
image_mode=ImageRefMode.PLACEHOLDER,
|
|
)
|
|
print(f" => produced {out_path / fname}.md")
|
|
|
|
res.document.save_as_html(
|
|
out_path / f"{fname}.html",
|
|
image_mode=ImageRefMode.EMBEDDED,
|
|
labels=[*DEFAULT_EXPORT_LABELS, DocItemLabel.FOOTNOTE],
|
|
split_page_view=True,
|
|
)
|
|
print(f" => produced {out_path / fname}.html")
|
|
|
|
pg_num = res.document.num_pages()
|
|
print("")
|
|
print(
|
|
f"Total document prediction time: {inference_time:.2f} seconds, pages: {pg_num}"
|
|
)
|
|
print("====================================================")
|
|
|
|
# return [source, f"{out_path / fname}.html", model_id, framework, inference_time, ]
|
|
return [source, model_id, framework, pg_num, inference_time, ]
|
|
|
|
if __name__ == "__main__":
|
|
|
|
sources = [
|
|
# "tests/data/2305.03393v1-pg9-img.png",
|
|
"tests/data/pdf/2305.03393v1-pg9.pdf",
|
|
]
|
|
|
|
out_path = Path("scratch")
|
|
out_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
## Use 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
|
|
pipeline_options.generate_page_images = True
|
|
|
|
## On GPU systems, enable flash_attention_2 with CUDA:
|
|
# pipeline_options.accelerator_options.device = AcceleratorDevice.CUDA
|
|
# pipeline_options.accelerator_options.cuda_use_flash_attention2 = True
|
|
|
|
rows = []
|
|
for vlm_options in [
|
|
# smoldocling_vlm_conversion_options, \
|
|
# smoldocling_vlm_mlx_conversion_options, \
|
|
# granite_vision_vlm_conversion_options, \
|
|
phi_vlm_conversion_options, \
|
|
# qwen25_vl_3b_vlm_mlx_conversion_options, \
|
|
# pixtral_12b_vlm_mlx_conversion_options,
|
|
# pixtral_12b_vlm_conversion_options,
|
|
]:
|
|
pipeline_options.vlm_options = vlm_options
|
|
|
|
## 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,
|
|
),
|
|
},
|
|
)
|
|
|
|
row = convert(sources=sources, converter=converter)
|
|
print("pipelines: \n", converter._get_initialized_pipelines())
|
|
|
|
rows.append(row)
|
|
|
|
print(tabulate(rows))
|
|
|
|
print("see if memory gets released ...")
|
|
time.sleep(10)
|