mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-12-08 20:58:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat!: Docling v2 (#117)

Browse Source 
---------

Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
Signed-off-by: Maxim Lysak <mly@zurich.ibm.com>
Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
Co-authored-by: Maxim Lysak <mly@zurich.ibm.com>
Co-authored-by: Michele Dolfi <dol@zurich.ibm.com>
Co-authored-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
This commit is contained in:
Christoph Auer
2024-10-16 21:02:03 +02:00
committed by GitHub
parent d504432c1e
commit 7d3be0edeb
144 changed files with 15180 additions and 3828 deletions
Download Patch File Download Diff File Expand all files Collapse all files

223
tests/verify_utils.py
View File

@@ -1,9 +1,10 @@
import json
import warnings
from pathlib import Path
from typing import List
from docling_core.types import BaseText
from docling_core.types import Document as DsDocument
from docling_core.types.doc import DoclingDocument
from pydantic import TypeAdapter
from pydantic.json import pydantic_encoder
@@ -79,32 +80,32 @@ def verify_cells(doc_pred_pages: List[Page], doc_true_pages: List[Page]):
return True
def verify_maintext(doc_pred: DsDocument, doc_true: DsDocument):
assert doc_true.main_text is not None, "doc_true cannot be None"
assert doc_pred.main_text is not None, "doc_true cannot be None"
assert len(doc_true.main_text) == len(
doc_pred.main_text
), f"document has different length of main-text than expected. {len(doc_true.main_text)}!={len(doc_pred.main_text)}"
for l, true_item in enumerate(doc_true.main_text):
pred_item = doc_pred.main_text[l]
# Validate type
assert (
true_item.obj_type == pred_item.obj_type
), f"Item[{l}] type does not match. expected[{true_item.obj_type}] != predicted [{pred_item.obj_type}]"
# Validate text ceels
if isinstance(true_item, BaseText):
assert isinstance(
pred_item, BaseText
), f"{pred_item} is not a BaseText element, but {true_item} is."
assert true_item.text == pred_item.text
return True
# def verify_maintext(doc_pred: DsDocument, doc_true: DsDocument):
# assert doc_true.main_text is not None, "doc_true cannot be None"
# assert doc_pred.main_text is not None, "doc_true cannot be None"
#
# assert len(doc_true.main_text) == len(
# doc_pred.main_text
# ), f"document has different length of main-text than expected. {len(doc_true.main_text)}!={len(doc_pred.main_text)}"
#
# for l, true_item in enumerate(doc_true.main_text):
# pred_item = doc_pred.main_text[l]
# # Validate type
# assert (
# true_item.obj_type == pred_item.obj_type
# ), f"Item[{l}] type does not match. expected[{true_item.obj_type}] != predicted [{pred_item.obj_type}]"
#
# # Validate text ceels
# if isinstance(true_item, BaseText):
# assert isinstance(
# pred_item, BaseText
# ), f"{pred_item} is not a BaseText element, but {true_item} is."
# assert true_item.text == pred_item.text
#
# return True
def verify_tables(doc_pred: DsDocument, doc_true: DsDocument, fuzzy: bool):
def verify_tables_v1(doc_pred: DsDocument, doc_true: DsDocument, fuzzy: bool):
if doc_true.tables is None:
# No tables to check
assert doc_pred.tables is None, "not expecting any table on this document"
@@ -130,6 +131,10 @@ def verify_tables(doc_pred: DsDocument, doc_true: DsDocument, fuzzy: bool):
assert true_item.data is not None, "documents are expected to have table data"
assert pred_item.data is not None, "documents are expected to have table data"
print("True: \n", true_item.export_to_dataframe().to_markdown())
print("Pred: \n", true_item.export_to_dataframe().to_markdown())
for i, row in enumerate(true_item.data):
for j, col in enumerate(true_item.data[i]):
@@ -137,7 +142,9 @@ def verify_tables(doc_pred: DsDocument, doc_true: DsDocument, fuzzy: bool):
# print("pred: ", pred_item.data[i][j].text)
# print("")
verify_text(true_item.data[i][j].text, pred_item.data[i][j].text, fuzzy)
verify_text(
true_item.data[i][j].text, pred_item.data[i][j].text, fuzzy=fuzzy
)
assert (
true_item.data[i][j].obj_type == pred_item.data[i][j].obj_type
@@ -146,14 +153,73 @@ def verify_tables(doc_pred: DsDocument, doc_true: DsDocument, fuzzy: bool):
return True
def verify_output(doc_pred: DsDocument, doc_true: DsDocument):
def verify_tables_v2(doc_pred: DoclingDocument, doc_true: DoclingDocument, fuzzy: bool):
if not len(doc_true.tables) > 0:
# No tables to check
assert len(doc_pred.tables) == 0, "not expecting any table on this document"
return True
else:
assert len(doc_pred.tables) > 0, "no tables predicted, but expected in doc_true"
assert verify_maintext(doc_pred, doc_true), "verify_maintext(doc_pred, doc_true)"
assert verify_tables(doc_pred, doc_true), "verify_tables(doc_pred, doc_true)"
# print("Expected number of tables: {}, result: {}".format(len(doc_true.tables), len(doc_pred.tables)))
assert len(doc_true.tables) == len(
doc_pred.tables
), "document has different count of tables than expected."
for l, true_item in enumerate(doc_true.tables):
pred_item = doc_pred.tables[l]
assert (
true_item.data.num_rows == pred_item.data.num_rows
), "table does not have the same #-rows"
assert (
true_item.data.num_cols == pred_item.data.num_cols
), "table does not have the same #-cols"
assert true_item.data is not None, "documents are expected to have table data"
assert pred_item.data is not None, "documents are expected to have table data"
print("True: \n", true_item.export_to_dataframe().to_markdown())
print("Pred: \n", true_item.export_to_dataframe().to_markdown())
for i, row in enumerate(true_item.data.grid):
for j, col in enumerate(true_item.data.grid[i]):
# print("true: ", true_item.data[i][j].text)
# print("pred: ", pred_item.data[i][j].text)
# print("")
verify_text(
true_item.data.grid[i][j].text,
pred_item.data.grid[i][j].text,
fuzzy=fuzzy,
)
assert (
true_item.data.grid[i][j].column_header
== pred_item.data.grid[i][j].column_header
), "table-cell should be a column_header but prediction isn't"
assert (
true_item.data.grid[i][j].row_header
== pred_item.data.grid[i][j].row_header
), "table-cell should be a row_header but prediction isn't"
assert (
true_item.data.grid[i][j].row_section
== pred_item.data.grid[i][j].row_section
), "table-cell should be a row_section but prediction isn't"
return True
# def verify_output(doc_pred: DsDocument, doc_true: DsDocument):
# #assert verify_maintext(doc_pred, doc_true), "verify_maintext(doc_pred, doc_true)"
# assert verify_tables_v1(doc_pred, doc_true), "verify_tables(doc_pred, doc_true)"
# return True
def verify_md(doc_pred_md: str, doc_true_md: str, fuzzy: bool):
return verify_text(doc_true_md, doc_pred_md, fuzzy)
@@ -162,7 +228,7 @@ def verify_dt(doc_pred_dt: str, doc_true_dt: str, fuzzy: bool):
return verify_text(doc_true_dt, doc_pred_dt, fuzzy)
def verify_conversion_result(
def verify_conversion_result_v1(
input_path: Path,
doc_result: ConversionResult,
generate: bool = False,
@@ -176,15 +242,18 @@ def verify_conversion_result(
), f"Doc {input_path} did not convert successfully."
doc_pred_pages: List[Page] = doc_result.pages
doc_pred: DsDocument = doc_result.output
doc_pred_md = doc_result.render_as_markdown()
doc_pred_dt = doc_result.render_as_doctags()
doc_pred: DsDocument = doc_result.legacy_document
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
doc_pred_md = doc_result.legacy_document.export_to_markdown()
doc_pred_dt = doc_result.legacy_document.export_to_document_tokens()
engine_suffix = "" if ocr_engine is None else f".{ocr_engine}"
pages_path = input_path.with_suffix(f"{engine_suffix}.pages.json")
json_path = input_path.with_suffix(f"{engine_suffix}.json")
md_path = input_path.with_suffix(f"{engine_suffix}.md")
dt_path = input_path.with_suffix(f"{engine_suffix}.doctags.txt")
gt_subpath = input_path.parent / "groundtruth" / "docling_v1" / input_path.name
pages_path = gt_subpath.with_suffix(f"{engine_suffix}.pages.json")
json_path = gt_subpath.with_suffix(f"{engine_suffix}.json")
md_path = gt_subpath.with_suffix(f"{engine_suffix}.md")
dt_path = gt_subpath.with_suffix(f"{engine_suffix}.doctags.txt")
if generate: # only used when re-generating truth
with open(pages_path, "w") as fw:
@@ -220,14 +289,86 @@ def verify_conversion_result(
# doc_pred, doc_true
# ), f"Mismatch in JSON prediction for {input_path}"
assert verify_tables(
doc_pred, doc_true, fuzzy
assert verify_tables_v1(
doc_pred, doc_true, fuzzy=fuzzy
), f"verify_tables(doc_pred, doc_true) mismatch for {input_path}"
assert verify_md(
doc_pred_md, doc_true_md, fuzzy
doc_pred_md, doc_true_md, fuzzy=fuzzy
), f"Mismatch in Markdown prediction for {input_path}"
assert verify_dt(
doc_pred_dt, doc_true_dt, fuzzy
doc_pred_dt, doc_true_dt, fuzzy=fuzzy
), f"Mismatch in DocTags prediction for {input_path}"
def verify_conversion_result_v2(
input_path: Path,
doc_result: ConversionResult,
generate: bool = False,
ocr_engine: str = None,
fuzzy: bool = False,
):
PageList = TypeAdapter(List[Page])
assert (
doc_result.status == ConversionStatus.SUCCESS
), f"Doc {input_path} did not convert successfully."
doc_pred_pages: List[Page] = doc_result.pages
doc_pred: DoclingDocument = doc_result.document
doc_pred_md = doc_result.document.export_to_markdown()
doc_pred_dt = doc_result.document.export_to_document_tokens()
engine_suffix = "" if ocr_engine is None else f".{ocr_engine}"
gt_subpath = input_path.parent / "groundtruth" / "docling_v2" / input_path.name
pages_path = gt_subpath.with_suffix(f"{engine_suffix}.pages.json")
json_path = gt_subpath.with_suffix(f"{engine_suffix}.json")
md_path = gt_subpath.with_suffix(f"{engine_suffix}.md")
dt_path = gt_subpath.with_suffix(f"{engine_suffix}.doctags.txt")
if generate: # only used when re-generating truth
with open(pages_path, "w") as fw:
fw.write(json.dumps(doc_pred_pages, default=pydantic_encoder))
with open(json_path, "w") as fw:
fw.write(json.dumps(doc_pred, default=pydantic_encoder))
with open(md_path, "w") as fw:
fw.write(doc_pred_md)
with open(dt_path, "w") as fw:
fw.write(doc_pred_dt)
else: # default branch in test
with open(pages_path, "r") as fr:
doc_true_pages = PageList.validate_json(fr.read())
with open(json_path, "r") as fr:
doc_true: DoclingDocument = DoclingDocument.model_validate_json(fr.read())
with open(md_path, "r") as fr:
doc_true_md = fr.read()
with open(dt_path, "r") as fr:
doc_true_dt = fr.read()
if not fuzzy:
assert verify_cells(
doc_pred_pages, doc_true_pages
), f"Mismatch in PDF cell prediction for {input_path}"
# assert verify_output(
# doc_pred, doc_true
# ), f"Mismatch in JSON prediction for {input_path}"
assert verify_tables_v2(
doc_pred, doc_true, fuzzy=fuzzy
), f"verify_tables(doc_pred, doc_true) mismatch for {input_path}"
assert verify_md(
doc_pred_md, doc_true_md, fuzzy=fuzzy
), f"Mismatch in Markdown prediction for {input_path}"
assert verify_dt(
doc_pred_dt, doc_true_dt, fuzzy=fuzzy
), f"Mismatch in DocTags prediction for {input_path}"