Repository: flairNLP/flair

Issue:
Different classification scores when using multi_class_prob
**Describe the bug**
When using the multi-class classifier's `predict` method with `multi_class_prob` argument set to `True` I would expect that the model behaves the same as with this flag set to `False` except that instead of probabilities for all classes it writes into the predicted labels only the label with highest score. However, that is not the case. 

In [L493](https://github.com/flairNLP/flair/blob/288f7ce4a0680ac8b3ab27d252f217d2bbd5a73c/flair/nn/model.py#L493) we check if the `multi_class_prob` is set, and if it is the scores are computed by pulling logits through the sigmoid in [L494](https://github.com/flairNLP/flair/blob/288f7ce4a0680ac8b3ab27d252f217d2bbd5a73c/flair/nn/model.py#L494), while if `multi_class_prob` is not set the scores are computed by pulling logits through the softmax in [L507](https://github.com/flairNLP/flair/blob/288f7ce4a0680ac8b3ab27d252f217d2bbd5a73c/flair/nn/model.py#L507). 

**To Reproduce**
Run the model that inherits from `DefaultClassifier` with and without `multi_class_prob` set and compare the scores for the highest scoring class. They scores are different. The sum of scores with `multi_class_prob` is not 1. 

**Expected behavior**
The scores should be the same. The sum of scores with `multi_class_prob` should be 1 (up to numerical error).

**Environment (please complete the following information):**
 - OS: MacOS (20.5.0 Darwin Kernel)
 - Version: 0.8.1 


Discussion:

Resolving patch:
diff --git a/flair/nn/model.py b/flair/nn/model.py
--- a/flair/nn/model.py
+++ b/flair/nn/model.py
@@ -13,7 +13,7 @@
 
 import flair
 from flair import file_utils
-from flair.data import DataPoint, Sentence, Dictionary, SpanLabel, Label
+from flair.data import DataPoint, Sentence, Dictionary, SpanLabel
 from flair.datasets import DataLoader, SentenceDataset
 from flair.training_utils import Result, store_embeddings
 
@@ -317,8 +317,9 @@ def evaluate(
 
         else:
             # issue error and default all evaluation numbers to 0.
-            log.error("ACHTUNG! No gold labels and no all_predicted_values found! Could be an error in your corpus or how you "
-                      "initialize the trainer!")
+            log.error(
+                "ACHTUNG! No gold labels and no all_predicted_values found! Could be an error in your corpus or how you "
+                "initialize the trainer!")
             accuracy_score = precision_score = recall_score = micro_f_score = macro_f_score = main_score = 0.
             classification_report = ""
             classification_report_dict = {}
@@ -410,7 +411,7 @@ def __init__(self,
     def multi_label_threshold(self):
         return self._multi_label_threshold
 
-    @multi_label_threshold.setter  
+    @multi_label_threshold.setter
     def multi_label_threshold(self, x):  # setter method
         if type(x) is dict:
             if 'default' in x:
@@ -419,7 +420,7 @@ def multi_label_threshold(self, x):  # setter method
                 raise Exception('multi_label_threshold dict should have a "default" key')
         else:
             self._multi_label_threshold = {'default': x}
-        
+
     def forward_loss(self, sentences: Union[List[DataPoint], DataPoint]) -> torch.tensor:
         scores, labels = self.forward_pass(sentences)
         return self._calculate_loss(scores, labels)
@@ -439,12 +440,11 @@ def _calculate_loss(self, scores, labels):
 
         return self.loss_function(scores, labels), len(labels)
 
-
     def predict(
             self,
             sentences: Union[List[Sentence], Sentence],
             mini_batch_size: int = 32,
-            multi_class_prob: bool = False,
+            return_probabilities_for_all_classes: bool = False,
             verbose: bool = False,
             label_name: Optional[str] = None,
             return_loss=False,
@@ -454,7 +454,7 @@ def predict(
         Predicts the class labels for the given sentences. The labels are directly added to the sentences.
         :param sentences: list of sentences
         :param mini_batch_size: mini batch size to use
-        :param multi_class_prob : return probability for all class for multiclass
+        :param return_probabilities_for_all_classes : return probabilities for all classes instead of only best predicted
         :param verbose: set to True to display a progress bar
         :param return_loss: set to True to return loss
         :param label_name: set this to change the name of the label type that is predicted
@@ -514,8 +514,7 @@ def predict(
 
                 # if anything could possibly be predicted
                 if len(label_candidates) > 0:
-
-                    if self.multi_label or multi_class_prob:
+                    if self.multi_label:
                         sigmoided = torch.sigmoid(scores)  # size: (n_sentences, n_classes)
                         n_labels = sigmoided.size(1)
                         for s_idx, (sentence, label) in enumerate(zip(sentences, label_candidates)):
@@ -524,20 +523,23 @@ def predict(
                                 if label_value == 'O': continue
                                 label_threshold = self._get_label_threshold(label_value)
                                 label_score = sigmoided[s_idx, l_idx].item()
-                                if label_score > label_threshold or multi_class_prob:
+                                if label_score > label_threshold or return_probabilities_for_all_classes:
                                     label.set_value(value=label_value, score=label_score)
                                     sentence.add_complex_label(label_name, copy.deepcopy(label))
-
                     else:
                         softmax = torch.nn.functional.softmax(scores, dim=-1)
-                        conf, idx = torch.max(softmax, dim=-1)
-
-                        for sentence, label, c, i in zip(sentences, label_candidates, conf, idx):
-                            label_value = self.label_dictionary.get_item_for_index(i.item())
-                            if label_value == 'O': continue
-                            label.set_value(value=label_value, score=c.item())
-
-                            sentence.add_complex_label(label_name, label)
+                        n_labels = softmax.size(1)
+                        for s_idx, (sentence, label) in enumerate(zip(sentences, label_candidates)):
+                            idx = torch.argmax(softmax[s_idx])
+                            for l_idx in range(n_labels):
+                                if not return_probabilities_for_all_classes and l_idx != idx:
+                                    continue
+                                else:
+                                    label_value = self.label_dictionary.get_item_for_index(l_idx)
+                                    if label_value == 'O': continue
+                                    label_score = softmax[s_idx, l_idx].item()
+                                    label.set_value(value=label_value, score=label_score)
+                                    sentence.add_complex_label(label_name, copy.deepcopy(label))
 
                 store_embeddings(batch, storage_mode=embedding_storage_mode)
 


Test patch:
diff --git a/tests/test_text_classifier.py b/tests/test_text_classifier.py
--- a/tests/test_text_classifier.py
+++ b/tests/test_text_classifier.py
@@ -130,7 +130,7 @@ def test_train_load_use_classifier_with_prob(results_base_path, tasks_base_path)
 
     sentence = Sentence("Berlin is a really nice city.")
 
-    model.predict(sentence, multi_class_prob=True)
+    model.predict(sentence, return_probabilities_for_all_classes=True)
 
     assert len(sentence.labels) > 1
 
@@ -145,9 +145,9 @@ def test_train_load_use_classifier_with_prob(results_base_path, tasks_base_path)
     sentence = Sentence("I love Berlin")
     sentence_empty = Sentence("       ")
 
-    loaded_model.predict(sentence, multi_class_prob=True)
-    loaded_model.predict([sentence, sentence_empty], multi_class_prob=True)
-    loaded_model.predict([sentence_empty], multi_class_prob=True)
+    loaded_model.predict(sentence, return_probabilities_for_all_classes=True)
+    loaded_model.predict([sentence, sentence_empty], return_probabilities_for_all_classes=True)
+    loaded_model.predict([sentence_empty], return_probabilities_for_all_classes=True)
 
     # clean up results directory
     shutil.rmtree(results_base_path)