公式增强

警告 本示例仅演示如何开发新的增强模型。它不运行实际的公式理解模型。

In [ ]

import logging
from collections.abc import Iterable
from pathlib import Path

import logging from collections.abc import Iterable from pathlib import Path

In [ ]

from docling_core.types.doc import DocItemLabel, DoclingDocument, NodeItem, TextItem

from docling_core.types.doc import DocItemLabel, DoclingDocument, NodeItem, TextItem

In [ ]

from docling.datamodel.base_models import InputFormat, ItemAndImageEnrichmentElement
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.models.base_model import BaseItemAndImageEnrichmentModel
from docling.pipeline.standard_pdf_pipeline import StandardPdfPipeline

from docling.datamodel.base_models import InputFormat, ItemAndImageEnrichmentElement from docling.datamodel.pipeline_options import PdfPipelineOptions from docling.document_converter import DocumentConverter, PdfFormatOption from docling.models.base_model import BaseItemAndImageEnrichmentModel from docling.pipeline.standard_pdf_pipeline import StandardPdfPipeline

In [ ]

class ExampleFormulaUnderstandingPipelineOptions(PdfPipelineOptions):
    do_formula_understanding: bool = True

class ExampleFormulaUnderstandingPipelineOptions(PdfPipelineOptions): do_formula_understanding: bool = True

In [ ]

# A new enrichment model using both the document element and its image as input
class ExampleFormulaUnderstandingEnrichmentModel(BaseItemAndImageEnrichmentModel):
    images_scale = 2.6

    def __init__(self, enabled: bool):
        self.enabled = enabled

    def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
        return (
            self.enabled
            and isinstance(element, TextItem)
            and element.label == DocItemLabel.FORMULA
        )

    def __call__(
        self,
        doc: DoclingDocument,
        element_batch: Iterable[ItemAndImageEnrichmentElement],
    ) -> Iterable[NodeItem]:
        if not self.enabled:
            return

        for enrich_element in element_batch:
            enrich_element.image.show()

            yield enrich_element.item

# 一个使用文档元素及其图像作为输入的新增强模型 class ExampleFormulaUnderstandingEnrichmentModel(BaseItemAndImageEnrichmentModel): images_scale = 2.6 def __init__(self, enabled: bool): self.enabled = enabled def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool: return ( self.enabled and isinstance(element, TextItem) and element.label == DocItemLabel.FORMULA ) def __call__( self, doc: DoclingDocument, element_batch: Iterable[ItemAndImageEnrichmentElement], ) -> Iterable[NodeItem]: if not self.enabled: return for enrich_element in element_batch: enrich_element.image.show() yield enrich_element.item

In [ ]

# How the pipeline can be extended.
class ExampleFormulaUnderstandingPipeline(StandardPdfPipeline):
    def __init__(self, pipeline_options: ExampleFormulaUnderstandingPipelineOptions):
        super().__init__(pipeline_options)
        self.pipeline_options: ExampleFormulaUnderstandingPipelineOptions

        self.enrichment_pipe = [
            ExampleFormulaUnderstandingEnrichmentModel(
                enabled=self.pipeline_options.do_formula_understanding
            )
        ]

        if self.pipeline_options.do_formula_understanding:
            self.keep_backend = True

    @classmethod
    def get_default_options(cls) -> ExampleFormulaUnderstandingPipelineOptions:
        return ExampleFormulaUnderstandingPipelineOptions()

# 如何扩展流水线 class ExampleFormulaUnderstandingPipeline(StandardPdfPipeline): def __init__(self, pipeline_options: ExampleFormulaUnderstandingPipelineOptions): super().__init__(pipeline_options) self.pipeline_options: ExampleFormulaUnderstandingPipelineOptions self.enrichment_pipe = [ ExampleFormulaUnderstandingEnrichmentModel( enabled=self.pipeline_options.do_formula_understanding ) ] if self.pipeline_options.do_formula_understanding: self.keep_backend = True @classmethod def get_default_options(cls) -> ExampleFormulaUnderstandingPipelineOptions: return ExampleFormulaUnderstandingPipelineOptions()

In [ ]

# Example main. In the final version, we simply have to set do_formula_understanding to true.
def main():
    logging.basicConfig(level=logging.INFO)

    input_doc_path = Path("./tests/data/pdf/2203.01017v2.pdf")

    pipeline_options = ExampleFormulaUnderstandingPipelineOptions()
    pipeline_options.do_formula_understanding = True

    doc_converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(
                pipeline_cls=ExampleFormulaUnderstandingPipeline,
                pipeline_options=pipeline_options,
            )
        }
    )
    doc_converter.convert(input_doc_path)

# 示例 main。在最终版本中，我们只需将 do_formula_understanding 设置为 true。 def main(): logging.basicConfig(level=logging.INFO) input_doc_path = Path("./tests/data/pdf/2203.01017v2.pdf") pipeline_options = ExampleFormulaUnderstandingPipelineOptions() pipeline_options.do_formula_understanding = True doc_converter = DocumentConverter( format_options={ InputFormat.PDF: PdfFormatOption( pipeline_cls=ExampleFormulaUnderstandingPipeline, pipeline_options=pipeline_options, ) } ) doc_converter.convert(input_doc_path)

In [ ]

if __name__ == "__main__":
    main()

if __name__ == "__main__": main()