mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-07-27 20:44:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added tokens/sec measurement, improved example

Browse Source 
Signed-off-by: Maksym Lysak <mly@zurich.ibm.com>
This commit is contained in:
Maksym Lysak 2025-01-15 10:22:48 +01:00
parent 437053572d
commit e0929781f4
2 changed files with 50 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
docling/models/smol_docling_model.py
View File

@ -63,7 +63,6 @@ class SmolDoclingModel(BasePageModel):
else: else:
with TimeRecorder(conv_res, "smolvlm"): with TimeRecorder(conv_res, "smolvlm"):
assert page.size is not None assert page.size is not None
start_time = time.time()
hi_res_image = page.get_image(scale=2.0) # 144dpi hi_res_image = page.get_image(scale=2.0) # 144dpi
# populate page_tags with predicted doc tags # populate page_tags with predicted doc tags
@ -95,19 +94,27 @@ class SmolDoclingModel(BasePageModel):
inputs = {k: v.to(self.device) for k, v in inputs.items()} inputs = {k: v.to(self.device) for k, v in inputs.items()}
prompt = prompt.replace("<end_of_utterance>", "") prompt = prompt.replace("<end_of_utterance>", "")
start_time = time.time()
# Call model to generate: # Call model to generate:
generated_ids = self.vlm_model.generate( generated_ids = self.vlm_model.generate(
**inputs, max_new_tokens=4096 **inputs, max_new_tokens=4096
) )
generation_time = time.time() - start_time
generated_texts = self.processor.batch_decode( generated_texts = self.processor.batch_decode(
generated_ids, skip_special_tokens=True generated_ids, skip_special_tokens=True
)[0] )[0]
num_tokens = len(generated_ids[0])
generated_texts = generated_texts.replace("Assistant: ", "") generated_texts = generated_texts.replace("Assistant: ", "")
page_tags = generated_texts page_tags = generated_texts
inference_time = time.time() - start_time inference_time = time.time() - start_time
tokens_per_second = num_tokens / generation_time
print("")
print(f"Page Inference Time: {inference_time:.2f} seconds") print(f"Page Inference Time: {inference_time:.2f} seconds")
print(f"Tokens/sec: {tokens_per_second:.2f}")
print("")
print("Page predictions:") print("Page predictions:")
print(page_tags) print(page_tags)

73
docs/examples/minimal_smol_docling.py
View File

@ -1,8 +1,11 @@
import json
import os import os
import time import time
from pathlib import Path from pathlib import Path
from urllib.parse import urlparse from urllib.parse import urlparse
import yaml
from docling.backend.docling_parse_backend import DoclingParseDocumentBackend from docling.backend.docling_parse_backend import DoclingParseDocumentBackend
from docling.datamodel.base_models import InputFormat from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions from docling.datamodel.pipeline_options import PdfPipelineOptions
@ -11,15 +14,16 @@ from docling.pipeline.vlm_pipeline import VlmPipeline
# source = "https://arxiv.org/pdf/2408.09869" # document per local path or URL # source = "https://arxiv.org/pdf/2408.09869" # document per local path or URL
# source = "tests/data/2305.03393v1-pg9-img.png" # source = "tests/data/2305.03393v1-pg9-img.png"
source = "tests/data/2305.03393v1-pg9.pdf" # source = "tests/data/2305.03393v1-pg9.pdf"
# source = "demo_data/page.png" # source = "demo_data/page.png"
# source = "demo_data/original_tables.pdf" # source = "demo_data/original_tables.pdf"
parsed = urlparse(source) sources = [
if parsed.scheme in ("http", "https"): "tests/data/2305.03393v1-pg9-img.png",
out_name = os.path.basename(parsed.path) # "tests/data/2305.03393v1-pg9.pdf",
else: # "demo_data/page.png",
out_name = os.path.basename(source) # "demo_data/original_tables.pdf",
]
pipeline_options = PdfPipelineOptions() pipeline_options = PdfPipelineOptions()
pipeline_options.generate_page_images = True pipeline_options.generate_page_images = True
@ -41,34 +45,41 @@ converter = DocumentConverter(
} }
) )
start_time = time.time() out_path = Path("scratch")
print("============") out_path.mkdir(parents=True, exist_ok=True)
print("starting...")
print("============")
print("")
result = converter.convert(source) for source in sources:
start_time = time.time()
print("================================================")
print("Processing... {}".format(source))
print("================================================")
print("")
print("------------") res = converter.convert(source)
print("MD:")
print("------------")
print("")
print(result.document.export_to_markdown())
Path("scratch").mkdir(parents=True, exist_ok=True) print("------------------------------------------------")
result.document.save_as_html( print("MD:")
filename=Path("scratch/{}.html".format(out_name)), print("------------------------------------------------")
image_mode=ImageRefMode.REFERENCED, print("")
labels=[*DEFAULT_EXPORT_LABELS, DocItemLabel.FOOTNOTE], print(res.document.export_to_markdown())
)
pg_num = result.document.num_pages() with (out_path / f"{res.input.file.stem}.html").open("w") as fp:
fp.write(res.document.export_to_html())
print("") with (out_path / f"{res.input.file.stem}.json").open("w") as fp:
inference_time = time.time() - start_time fp.write(json.dumps(res.document.export_to_dict()))
print(f"Total document prediction time: {inference_time:.2f} seconds, pages: {pg_num}")
print("============") with (out_path / f"{res.input.file.stem}.yaml").open("w") as fp:
fp.write(yaml.safe_dump(res.document.export_to_dict()))
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}"
)
print("================================================")
print("done!") print("done!")
print("============") print("================================================")
# output: ## Docling Technical Report [...]"