Callbacks between Pipeline Steps

[simonschoe]

I am wondering if there is an option to apply callbacks after each pipeline step and, if not, do have any suggestions on how to implement it?

Exemplary scenario: sending status updates (e.g. about progress) after each step.

[dolfim-ibm]

You could implement your custom pipeline and in self.build_pipe you add extra steps.

Some references

• The PDF pipeline at https://github.com/DS4SD/docling/blob/main/docling/pipeline/standard_pdf_pipeline.py#L66
• Example extending the PDF pipeline with new enrichment at https://ds4sd.github.io/docling/examples/develop_picture_enrichment/

[simonschoe]

Thanks for the pointers! I have found this implementation to work while requiring little duplicate code:

def custom_callback(m):
    # implement logic
    pass
    
class PdfPipelineWithCallbacks(StandardPdfPipeline):

    def __init__(self, pipeline_options: PdfPipelineOptions, callback=custom_callback):
        super().__init__(pipeline_options)
        self.pipeline_options: PdfPipelineOptions
        self.callback = callback

        self.build_pipe = [
            self._wrap_with_callback(model) for model in self.build_pipe
        ]

    def _wrap_with_callback(self, model):
        def wrapped_model(*args, **kwargs):
            if self.callback:
                self.callback(model)
            return model(*args, **kwargs)
        return wrapped_model

doc_converter = DocumentConverter(
    format_options={
        InputFormat.PDF: PdfFormatOption(
            pipeline_cls=PdfPipelineWithCallbacks,
            pipeline_options=pipeline_options,
            backend=DoclingParseV2DocumentBackend,
        )
    }
)

Do you see any obvious errors in this approach?

EDIT: @dolfim-ibm I just realized that this callback would be applied after each batch (i.e., each time one of the pipeline models is run). I see that elements_batch_size: int = 16. Any chance to know what the total batch size is?

[cau-git]

@simonschoe Great share! The batch size for the page batches is controlled here, and the default is 4.