diff --git a/docling/utils/nemotron_parse_utils.py b/docling/utils/nemotron_parse_utils.py index 134f794201..2a4ac1c15a 100644 --- a/docling/utils/nemotron_parse_utils.py +++ b/docling/utils/nemotron_parse_utils.py @@ -31,6 +31,7 @@ from docling.backend.md_backend import MarkdownDocumentBackend from docling.datamodel.base_models import InputFormat from docling.datamodel.document import InputDocument +from docling.utils.vlm_utils import strip_math_delimiters _log = logging.getLogger(__name__) @@ -321,7 +322,7 @@ def parse_nemotron_parse_v2( else: label = _TEXT_LABELS.get(region.label, DocItemLabel.TEXT) text = ( - raw_text + strip_math_delimiters(raw_text) if label == DocItemLabel.FORMULA else _markdown_to_text(raw_text) ) diff --git a/docling/utils/vlm_utils.py b/docling/utils/vlm_utils.py index d3ed9c5d11..1c8f0ed42e 100644 --- a/docling/utils/vlm_utils.py +++ b/docling/utils/vlm_utils.py @@ -71,6 +71,19 @@ def parse_markdown_heading(text: str) -> tuple[str, int | None]: return match.group(2), len(match.group(1)) +def strip_math_delimiters(text: str) -> str: + """Remove one matching outer math pair without interpreting its LaTeX.""" + text = text.strip() + for opening, closing in (("$$", "$$"), (r"\[", r"\]"), (r"\(", r"\)"), ("$", "$")): + if ( + len(text) >= len(opening) + len(closing) + and text.startswith(opening) + and text.endswith(closing) + ): + return text[len(opening) : -len(closing)].strip() + return text + + def strip_stop_strings(texts: list[str], stop_strings: list[str]) -> list[str]: """Strip stop strings from decoded texts. diff --git a/tach.toml b/tach.toml index b8635b2581..8e2d8f2c2a 100644 --- a/tach.toml +++ b/tach.toml @@ -185,7 +185,7 @@ depends_on = [] [[modules]] path = "docling.utils.nemotron_parse_utils" layer = "core" -depends_on = ["docling.backend", "docling.datamodel"] +depends_on = ["docling.backend", "docling.datamodel", "docling.utils"] [[modules]] path = "docling.utils" diff --git a/tests/test_nemotron_formula_exports.py b/tests/test_nemotron_formula_exports.py new file mode 100644 index 0000000000..da994ba0e0 --- /dev/null +++ b/tests/test_nemotron_formula_exports.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: The Docling Contributors +# SPDX-License-Identifier: MIT + +import pytest +from docling_core.types.doc import FormulaItem, Size + +from docling.utils.nemotron_parse_utils import parse_nemotron_parse_v2 + + +@pytest.mark.parametrize( + "opening, closing", + [("$", "$"), ("$$", "$$"), (r"\(", r"\)"), (r"\[", r"\]"), ("", "")], +) +def test_formula_export_has_one_delimiter_pair(opening: str, closing: str) -> None: + latex = r"\left\{x_1 * y_2 : x > 0\right\}" + raw = f"{opening}{latex}{closing}" + source = f"{raw}" + size = Size(width=600, height=800) + doc = parse_nemotron_parse_v2(source, size, size, page_no=1) + formula = doc.texts[0] + assert isinstance(formula, FormulaItem) + assert formula.text == latex + assert formula.orig == raw + assert formula.prov + assert doc.export_to_markdown() == f"$${latex}$$" + + +def test_multiline_formula_keeps_latex_and_line_breaks() -> None: + latex = "\\begin{aligned}\nx &= y \\\\\nz &= w\n\\end{aligned}" + source = f"$$\n{latex}\n$$" + size = Size(width=600, height=800) + doc = parse_nemotron_parse_v2(source, size, size, page_no=1) + assert doc.texts[0].text == latex + assert doc.export_to_markdown() == f"$${latex}$$" diff --git a/tests/test_nemotron_parse_v2.py b/tests/test_nemotron_parse_v2.py index 3ff1941cbd..6bb03c7db8 100644 --- a/tests/test_nemotron_parse_v2.py +++ b/tests/test_nemotron_parse_v2.py @@ -183,7 +183,7 @@ def test_parse_nemotron_markdown_and_semantic_classes() -> None: item.text for item in document.texts if item.label == DocItemLabel.LIST_ITEM ] == ["First item", "Second item"] assert any( - item.label == DocItemLabel.FORMULA and item.text == r"\(x^2\)" + item.label == DocItemLabel.FORMULA and item.text == r"x^2" for item in document.texts ) assert len(document.pictures) == 2