ebanalyse / NERDA

Framework for fine-tuning pretrained transformers for Named-Entity Recognition (NER) tasks
MIT License
153 stars 35 forks source link

Extracting softmax scores on prediction #32

Open AndersGiovanni opened 2 years ago

AndersGiovanni commented 2 years ago

Hi @smaakage85,

Thank you very much for your work on the NERDA project, very much appreciated!

In the predict() function found in predictions.py you've added the functionality to return return_confidence which is super nice! However, when you extract the most likely tag using values, indices = outputs[i].max(dim=1) (from line 104) the prediction scores have not been put in a softmax function which would give confidence scores of each tag relative to the others.

It would be very nice to either

  1. Being able to get all the tag probabilities for each prediction.
  2. Apply a softmax function in either predict() or in forward() (NERDANetwork) like they do in this basic pytorch tutorial.