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/

[cau-git]

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

[simonschoe]

@cau-git @dolfim-ibm I tried to use the same logic to wrap a callback around an enrichment model:

class PdfPipelineWithCallback(StandardPdfPipeline):

    def __init__(self, pipeline_options: PdfPipelineOptions, callback: callable = None) -> None:
        
        self.build_pipe = [
            self._wrap_with_callback(model) for model in self.build_pipe
        ]
        self.enrichment_pipe = [
            self._wrap_with_callback(model) for model in self.enrichment_pipe
        ]

    def _wrap_with_callback(self, model: Any) -> callable:
        def wrapped_model(*args: Any, **kwargs: Any) -> Any:
            ...
            print(str(model))
            return model(*args, **kwargs)
        return wrapped_model

However, when doing so, I get the following error:
AttributeError: 'function' object has no attribute 'prepare_element'

Could you help me pin down the reason? :)

[dolfim-ibm]

I think it is because the base model of the enrichment_pipe has more functions, not only the __call__

[simonschoe]

I think it is because the base model of the enrichment_pipe has more functions, not only the __call__

Makes sense! Do you have any ideas how I could adjust the code above to wrap an enrichment model?

[cau-git]

@simonschoe Here is a complete example for an enrichment model: https://github.com/DS4SD/docling/blob/main/docs/examples/develop_formula_understanding.py

[ccchan-cpii]

@simonschoe can you post the completed code here with the enrichment model?

[simonschoe]

@ccchan-cpii Not sure if this is the best possible implementation of a wrapper, but it could look somethin like this (feedback appreciated!):

class ModelWrapper:
    """Generic wrapper around Docling model to allow for callback."""

    def __init__(
        self,
        model: Callable[..., object],
        callback: Callable,
    ) -> None:
        """Initialize CallableModelWrapper."""
        self._model = model
        self._callback = callback

        self._batch_id = 0
        self._page_per_batch = BatchConcurrencySettings().page_batch_size
        self._elements_batch_size = BatchConcurrencySettings().elements_batch_size
        self._total_batches = None
        self._total_pictures = None

    def __call__(
        self,
        *args: tuple[ConversionResult, Iterable[Any]],
        **kwargs: dict,
    ) -> object:
        """Wrap __call__ method to allow for callback after __call__ is executed."""
        result = self._model(*args, **kwargs)

        if isinstance(self._model, PageAssembleModel):

            if not self._total_batches and args:
                self._total_batches = math.ceil(
                    min(conv_result.input.limits.page_range[1], conv_result.input.page_count) / self._page_per_batch,
                )
            self._batch_id += 1

            # callback logic
            self._callback(self._batch_id, self._total_batches)

        if isinstance(self._model, DocumentPictureClassifier):

            if not self._total_pictures and kwargs:
                self._total_pictures = math.ceil(
                    len(kwargs["doc"].pictures) / self._elements_batch_size,
                )
            self._batch_id += 1

            # callback logic
            self._callback(self._batch_id, self._total_pictures)

        return result

    def __getattr__(self, attr: str) -> object:
        """Delegate attribute access to the wrapped model instance for all other methods."""
        return getattr(self._model, attr)



class PdfPipelineWithCallback(StandardPdfPipeline):
    """Wrapper around StandardPdfPipeline to allow for model callback."""

    def __init__(self, pipeline_options: PdfPipelineOptions, callback: callable) -> None:
        """Initialize PdfPipelineWithCallback."""
        super().__init__(pipeline_options)
        self._callback = callback

        self.build_pipe: list[ModelWrapper] = [
            ModelWrapper(model, self._callback) for model in self.build_pipe
        ]
        self.enrichment_pipe: list[ModelWrapper] = [
            ModelWrapper(model, self._callback) for model in self.enrichment_pipe
        ]