Repository: flairNLP/flair

Issue:
Allow per-label threshold for multi-label classifiers
Currently there is a single threshold for multi-label classifiers, meaning that the same threshold is applied to all labels. It is sometimes beneficial to have different thresholds for different labels. 

Model class has `multi_label_threshold` member. Currently it's a float. We could make it dictionary:
```
model.multi_label_threshold = {'label_1': label_1_threshold, ..., 'default': default_threshold}
```
so when the label is explicitly mentioned in the dictionary, the corresponding threshold is applied, otherwise `default_threshold` is applied.


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
@@ -381,6 +381,20 @@ def __init__(self,
         else:
             self.loss_function = torch.nn.CrossEntropyLoss(weight=self.loss_weights)
 
+    @property
+    def multi_label_threshold(self):
+        return self._multi_label_threshold
+
+    @multi_label_threshold.setter  
+    def multi_label_threshold(self, x):  # setter method
+        if type(x) is dict:
+            if 'default' in x:
+                self._multi_label_threshold = x
+            else:
+                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)
@@ -400,6 +414,7 @@ def _calculate_loss(self, scores, labels):
 
         return self.loss_function(scores, labels), len(labels)
 
+
     def predict(
             self,
             sentences: Union[List[Sentence], Sentence],
@@ -476,17 +491,17 @@ def predict(
                 if len(label_candidates) > 0:
 
                     if self.multi_label or multi_class_prob:
-                        sigmoided = torch.sigmoid(scores)
-                        s_idx = 0
-                        for sentence, label in zip(sentences, label_candidates):
-                            for idx in range(sigmoided.size(1)):
-                                if sigmoided[s_idx, idx] > self.multi_label_threshold or multi_class_prob:
-                                    label_value = self.label_dictionary.get_item_for_index(idx)
-                                    if label_value == 'O': continue
-                                    label.set_value(value=label_value, score=sigmoided[s_idx, idx].item())
+                        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)):
+                            for l_idx in range(n_labels):
+                                label_value = self.label_dictionary.get_item_for_index(l_idx)
+                                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:
+                                    label.set_value(value=label_value, score=label_score)
                                     sentence.add_complex_label(label_name, copy.deepcopy(label))
-                            s_idx += 1
-
                     else:
                         softmax = torch.nn.functional.softmax(scores, dim=-1)
                         conf, idx = torch.max(softmax, dim=-1)
@@ -503,6 +518,13 @@ def predict(
             if return_loss:
                 return overall_loss, label_count
 
+    def _get_label_threshold(self, label_value):
+        label_threshold = self.multi_label_threshold['default']
+        if label_value in self.multi_label_threshold:
+            label_threshold = self.multi_label_threshold[label_value]
+
+        return label_threshold
+
     def _obtain_labels(
             self, scores: List[List[float]], predict_prob: bool = False
     ) -> List[List[Label]]:
@@ -526,9 +548,11 @@ def _get_multi_label(self, label_scores) -> List[Label]:
 
         results = list(map(lambda x: sigmoid(x), label_scores))
         for idx, conf in enumerate(results):
-            if conf > self.multi_label_threshold:
-                label = self.label_dictionary.get_item_for_index(idx)
-                labels.append(Label(label, conf.item()))
+            label_value = self.label_dictionary.get_item_for_index(idx)
+            label_threshold = self._get_label_threshold(label_value)
+            label_score = conf.item()
+            if label_score > label_threshold:
+                labels.append(Label(label_value, label_score))
 
         return labels
 


Test patch: