add namedtuple preds output
# https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/examples/run_classifier.py
import random
+from collections import namedtuple
import numpy as np
import torch
@@ -168,11 +169,15 @@ class BERTSequenceClassifier:
If None is specified, all available GPUs
will be used. Defaults to None.
batch_size (int, optional): Scoring batch size. Defaults to 32.
- probabilities: If true, the predicted probability distribution
- is returned; otherwise, the predicted classes are returned.
+ class_proba (bool, str, optional):
+ If True, the predicted probability distribution is returned;
+ If False, the predicted classes are returned;
+ If "both", both the predicted classes and probabilities
+ are returned.
Defaults to False.
Returns:
- [ndarray]: Predicted classes.
+ ndarray, namedtuple(ndarray, ndarray): Predicted classes,
+ probabilities, or both.
"""
device = get_device("cpu" if num_gpus == 0 else "gpu")
@@ -204,7 +209,12 @@ class BERTSequenceClassifier:
preds = np.concatenate(preds)
- if probabilities:
+ if class_proba == "both":
+ return namedtuple("Predictions", "probabilities classes")(
+ nn.Softmax(dim=1)(torch.Tensor(preds)).numpy(),
+ preds.argmax(axis=1),
+ )
+ elif class_proba:
return nn.Softmax(dim=1)(torch.Tensor(preds)).numpy()
else:
return preds.argmax(axis=1)
