How to compute a squad metric?

[saneim]

Hallo,
I am trying to train a BERT model with the squad dataset.
Since the model has a maximum input length, examples with longer contexts must be split into several samples.
In the validation_epoch_end() and test_epoch_end() methods I want to compute the squad metric (https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/text/squad.py) that needs the predicted answers in text format.
But for each sample the model predicts only the start and end token.
To generate the answers in text format I need the dataset from the val/test_dataloader as well as the original val/test_dataset (with text instead of input_ids, offset_mappings etc.)
So I was wondering how to access the original dataset in those methods.
Of course I can use self.trainer.datamodule.original_val_dataset but then I need to ensure that the trainer uses a datamodule and not only dataloaders. What is the best practice for this case?
And how do I use this model in production? Can I compute the text answers in the LightningModule and return them when predict() is called ?
Or should predict() only return the start and end tokens and the text answer is computed outside of the LightningModule?
Thank you in advance!

[rohitgr7]

I'd suggest accessing the dataloader/dataset isn't a reliable solution, but rather you can return the original text and offests in Dataset.getitem or you can store the offests and text in the state variables for easy access.

Here are some examples I worked with on kaggle:
For the first one where you return the offests/text: https://www.kaggle.com/rohitgr/roberta-with-pytorch-lightning-train-test-lb-0-710
This one uses an old version of lightning, but should still be relevant I guess.

Here I store them in my datamodule: https://www.kaggle.com/rohitgr/chaii-q-a-with-pytorch-lightining
this one is recent so should be working.

[saneim]

Thank you very much for your reply. I did a bit more research and want to ask if my understanding is correct:
The validation/test/train/predict_step and validation/test/train/predict_epoch_end methods are only for developing and testing the model? In production the model is exported and the forward method is used.
So the post processing for production (convert predictions into answers) must be done outside of the LightningModule.
Since I assume that I have start and end positions for each sample during training, I could also assume that the DataModule gives me the original data sets during development. So I could check in the epoch_end methods if the trainer has a DataModule. If yes I can use the squad-metric otherwise I could just calculate the loss.
Would this be a good solution or is this quick and dirty?

[rohitgr7]

if you are calculating loss/metric then you should not do the postprocessing inside the forward method because as you said correctly that forward is used only for inferencing assuming you don't have actual predictions. For development, you can do it inside *_step/*_epoch_end methods but during inference, it can be done outside LightningModule.