Repository: flairNLP/flair

Issue:
[Text Classification] Fix bug in single label training
When training a text classifier with a single label task, the classifier predicts just one class. 


Discussion:

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -437,21 +437,16 @@ def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
 
 class DocumentMeanEmbeddings(DocumentEmbeddings):
 
-    def __init__(self, word_embeddings: List[TokenEmbeddings], reproject_words: bool = True):
+    def __init__(self, word_embeddings: List[TokenEmbeddings]):
         """The constructor takes a list of embeddings to be combined."""
         super().__init__()
 
         self.embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=word_embeddings)
         self.name: str = 'document_mean'
-        self.reproject_words: bool = reproject_words
-        self.static_embeddings: bool = not reproject_words
 
         self.__embedding_length: int = 0
         self.__embedding_length = self.embeddings.embedding_length
 
-        self.word_reprojection_map = torch.nn.Linear(self.__embedding_length, self.__embedding_length)
-        torch.nn.init.xavier_uniform_(self.word_reprojection_map.weight)
-
         if torch.cuda.is_available():
             self.cuda()
 
@@ -472,7 +467,7 @@ def embed(self, sentences: Union[List[Sentence], Sentence]):
         for sentence in sentences:
             if self.name not in sentence._embeddings.keys(): everything_embedded = False
 
-        if not everything_embedded or not self.static_embeddings:
+        if not everything_embedded:
 
             self.embeddings.embed(sentences)
 
@@ -486,9 +481,6 @@ def embed(self, sentences: Union[List[Sentence], Sentence]):
                 if torch.cuda.is_available():
                     word_embeddings = word_embeddings.cuda()
 
-                if self.reproject_words:
-                    word_embeddings = self.word_reprojection_map(word_embeddings)
-
                 mean_embedding = torch.mean(word_embeddings, 0)
 
                 sentence.set_embedding(self.name, mean_embedding.unsqueeze(0))
diff --git a/flair/models/text_classification_model.py b/flair/models/text_classification_model.py
--- a/flair/models/text_classification_model.py
+++ b/flair/models/text_classification_model.py
@@ -156,7 +156,7 @@ def _get_multi_label(self, label_scores) -> List[Label]:
         return labels
 
     def _get_single_label(self, label_scores) -> List[Label]:
-        conf, idx = torch.max(label_scores[0], 0)
+        conf, idx = torch.max(label_scores, 0)
         label = self.label_dictionary.get_item_for_index(idx.item())
 
         return [Label(label, conf.item())]


Test patch:
diff --git a/tests/test_text_classifier_trainer.py b/tests/test_text_classifier_trainer.py
--- a/tests/test_text_classifier_trainer.py
+++ b/tests/test_text_classifier_trainer.py
@@ -37,7 +37,7 @@ def test_text_classifier_mulit_label():
     label_dict = corpus.make_label_dictionary()
 
     glove_embedding: WordEmbeddings = WordEmbeddings('en-glove')
-    document_embeddings: DocumentMeanEmbeddings = DocumentMeanEmbeddings([glove_embedding], True)
+    document_embeddings: DocumentMeanEmbeddings = DocumentMeanEmbeddings([glove_embedding])
 
     model = TextClassifier(document_embeddings, label_dict, True)