Looking for tool class to do predictions

PhilipMay
Wed, 07 Oct 2020 09:45:13 +0000

Hi,
I am looking for a tool class to do predictions with BERT models. Like the Trainer class just for prediction time. There is much more to do than just doing logits = model(input_ids). When you have large input lists of text you have to do this in batches or you get an OOM from your GPU and there is some more advanced stuff like smart batching (sorting the sentences by token length and batch short sentences with other short sentences to have short batches with minimal padding).

Is there something like this?

Thanks
Philip