Looking for tool class to do predictions

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?


Hi Philip,

maybe the HF Pipelines implementation could help:

It should also support e.g. multiple sentences (List[str]) as inputs and you could use of course own trained and fine-tuned models. Detailed implementation can be found here:

To use it for text classification/sentiment analysis:

nlp_sentence_classif = pipeline('sentiment-analysis')
nlp_sentence_classif('Such a nice weather outside !')

and it also outputs the score:

[{'label': 'POSITIVE', 'score': 0.9997656}]

I like using Pipelines especially for testing NER models :hugs:

I bet! :wink:

Thanks for the answer.

@stefan-it although I see no support for batching. At least not in the TextClassificationPipeline.

When I follow the naive approach:

model = AutoModelForSequenceClassification.from_pretrained(model_dir)
tokenizer = AutoTokenizer.from_pretrained(model_dir)
pipeline = TextClassificationPipeline(

I am getting an CUDA OOM if the list (unlabeled_text_list) is long…