Fix torch.cat() issue when processing large number of documents with TransformersModelForTokenClassificationNerStep

[paluchasz]

The Issue

Noticed a weird problem occurring in the evaluation script when trying to naively process a large number (365) of Kazu documents with the TransformersModelForTokenClassificationNerStep step using the MPS device.

The 365 documents used totalled over 14k sections and were being processed with a newly trained 400MB model. Performing this on a Mac M3 with MPS, I saw Python's memory usage peak at 18GB:

In the end the step failed to predict any entities but without any exceptions thrown. The result was a weird phenomenon inside https://github.com/AstraZeneca/KAZU/blob/main/kazu/steps/ner/hf_token_classification.py:

  def get_multilabel_activations(self, loader: DataLoader) -> Tensor:
        """Get a tensor consisting of confidences for labels in a multi label
        classification context.

        :param loader:
        :return:
        """
        with torch.no_grad():
            results = torch.cat(
                tuple(self.model(**batch.to(self.device)).logits for batch in loader)
            ).to(self.device)
        return results.heaviside(torch.tensor([0.0]).to(self.device)).int().to("cpu")

Where torch.cat was producing a tensor full of zeros, indicating the model has not found any entities. This is likely due to torch.cat exceeding the allocated memory of the device.

The Fix

The fix is in two places. Firstly in the evaluate script we now process the documents in batches through the pipeline. However, to stop a user naively processing many documents with the Kazu pipeline and hitting this issue, there is also a fix inside the TransformersModelForTokenClassificationNerStep. This offloads the model logits onto CPU before concatenation.

Testing Performance

Here we perform the test with the naive call to the step with all the documents at once as before and test the version of TransformersModelForTokenClassificationNerStep before and after the change.

Before the change we observe a peak memory usage of 18GB and it takes 690s to process all the documents. With the new implementation we see a peak memory usage of 4GB and it takes 680s to process all the documents - also fixing the weird issue. Thus there doesn't seem to be any performance degradation in executing torch.cat on cpu vs mps. Cuda device was not tested however.

General Test for single label classification

A test script with the default model pipeline and Kazu model pack was run as a sanity check. The integration tests will now also be run.

Note

There is also a small refactor moving some functions from train_multilabel_ner to modelling_utils. Individual changes can be seen at commit level.

[mariosaenger]

LGTM. I only added minor comments

[mariosaenger]

LGTM