Repository: flairNLP/flair

Issue:
Introducing Attentional BiLSTM for Text Classification
As per my understanding ,  Flair by default uses the **standard BiLSTM model** for text classification.

I guess introducing **Attentional BiLSTM** might improve this maybe?


An easy solution to this might be to just introduce a parameter named **attention** in the **TextClassifier** class that decides whether to use attention or not 


Discussion:
If this issue is a go, I would love to take up the challenge
That's a great idea - we'd very much appreciate it if you added this! :)
@alanakbik , I am on it
@53X excited to see this implementation

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -1505,7 +1505,7 @@ def __init__(self,
                  rnn_layers=1,
                  reproject_words: bool = True,
                  reproject_words_dimension: int = None,
-                 bidirectional: bool = False,
+                 bidirectional: bool = True,
                  dropout: float = 0.5,
                  word_dropout: float = 0.0,
                  locked_dropout: float = 0.0):
@@ -1524,19 +1524,19 @@ def __init__(self,
         """
         super().__init__()
 
-        self.embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embeddings)
+        self.embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embeddings) 
 
         self.reproject_words = reproject_words
         self.bidirectional = bidirectional
-
-        self.length_of_all_token_embeddings: int = self.embeddings.embedding_length
+        
+        self.length_of_all_token_embeddings: int = self.embeddings.embedding_length 
 
         self.name = 'document_lstm'
         self.static_embeddings = False
 
         self.__embedding_length: int = hidden_size
         if self.bidirectional:
-            self.__embedding_length *= 4
+            self.__embedding_length *= 4 
 
         self.embeddings_dimension: int = self.length_of_all_token_embeddings
         if self.reproject_words and reproject_words_dimension is not None:
@@ -1563,19 +1563,19 @@ def __init__(self,
         self.to(flair.device)
 
     @property
-    def embedding_length(self) -> int:
+    def embedding_length(self) -> int:   
         return self.__embedding_length
 
-    def embed(self, sentences: Union[List[Sentence], Sentence]):
+    def embed(self, sentences: Union[List[Sentence], Sentence], return_sequences: bool = False):
         """Add embeddings to all sentences in the given list of sentences. If embeddings are already added, update
          only if embeddings are non-static."""
 
-        if type(sentences) is Sentence:
+        if type(sentences) is Sentence:   
             sentences = [sentences]
 
         self.rnn.zero_grad()
 
-        sentences.sort(key=lambda x: len(x), reverse=True)
+        sentences.sort(key=lambda x: len(x), reverse=True) 
 
         self.embeddings.embed(sentences)
 
@@ -1640,17 +1640,25 @@ def embed(self, sentences: Union[List[Sentence], Sentence]):
         # --------------------------------------------------------------------
         # EXTRACT EMBEDDINGS FROM LSTM
         # --------------------------------------------------------------------
-        for sentence_no, length in enumerate(lengths):
-            last_rep = outputs[length - 1, sentence_no]
 
-            embedding = last_rep
-            if self.bidirectional:
-                first_rep = outputs[0, sentence_no]
-                embedding = torch.cat([first_rep, last_rep], 0)
+        if return_sequences:
+            
+            return outputs
+        
+        else:
+            
+            for sentence_no, length in enumerate(lengths):
+                last_rep = outputs[length - 1, sentence_no]
+
+                embedding = last_rep
+                if self.bidirectional:
+                    first_rep = outputs[0, sentence_no]
+                    embedding = torch.cat([first_rep, last_rep], 0)
 
-            sentence = sentences[sentence_no]
-            sentence.set_embedding(self.name, embedding)
+                sentence = sentences[sentence_no]
+                sentence.set_embedding(self.name, embedding)
 
+           
     def _add_embeddings_internal(self, sentences: List[Sentence]):
         pass
 
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
@@ -11,10 +11,48 @@
 from flair.data import Dictionary, Sentence, Label
 from flair.file_utils import cached_path
 from flair.training_utils import convert_labels_to_one_hot, clear_embeddings
-
+import torch.nn.functional as F
 
 log = logging.getLogger('flair')
 
+class Attention(flair.nn.Model):
+
+    '''
+    This class implements soft-attention
+    '''
+    def __init__(self, batch_first: bool = True):
+
+        super(Attention, self).__init__()
+        self.batch_first: bool = batch_first
+    
+    def forward(self, attention_candidates: torch.Tensor, attention_size: int, weighted_sum_candidates: torch.Tensor = None) -> torch.Tensor:
+
+        if self.batch_first:
+            
+            self.attention_tensor: torch.Tensor = torch.empty(attention_candidates.size(0), attention_size, 1) 
+            torch.nn.init.ones_(self.attention_tensor)
+            weights = F.softmax(torch.matmul(attention_candidates, self.attention_tensor), dim=1)
+        
+        else:
+            
+            dummy = torch.empty(attention_candidates.size(1), attention_candidates.size(0), attention_candidates.size(2))
+            self.attention_tensor: torch.Tensor = torch.empty(attention_candidates.size(1), attention_size, 1)
+            torch.nn.init.ones_(self.attention_tensor)    
+            for i in range(attention_candidates.size(1)):
+                dummy[i] = torch.cat([attention_candidates[j, i, :] for j in range(attention_candidates.size(0))], dim=0).view(attention_candidates.size(0), attention_candidates.size(2))
+            weights = F.softmax(torch.matmul(dummy, self.attention_tensor), dim=1)  
+
+        if weighted_sum_candidates is None:
+            if self.batch_first:
+                weighting = torch.sum(weights*attention_candidates, dim=1, keepdim=True)
+            else:
+                weighting = torch.sum(weights*dummy, dim=1, keepdim=True)
+
+        else:
+            weighting = torch.sum(weights*weighted_sum_candidates, dim=1, keepdim=True) 
+
+        return weighting
+
 
 class TextClassifier(flair.nn.Model):
     """
@@ -27,13 +65,21 @@ class TextClassifier(flair.nn.Model):
     def __init__(self,
                  document_embeddings: flair.embeddings.DocumentEmbeddings,
                  label_dictionary: Dictionary,
-                 multi_label: bool):
+                 multi_label: bool, 
+                 attention: bool = False):
 
         super(TextClassifier, self).__init__()
 
         self.document_embeddings = document_embeddings
         self.label_dictionary: Dictionary = label_dictionary
         self.multi_label = multi_label
+        self.attention = attention
+
+        if self.attention:
+
+            self.attention_model = Attention(batch_first=False)
+            print('USING ATTENTION')
+
 
         self.document_embeddings: flair.embeddings.DocumentLSTMEmbeddings = document_embeddings
 
@@ -53,12 +99,18 @@ def _init_weights(self):
         nn.init.xavier_uniform_(self.decoder.weight)
 
     def forward(self, sentences) -> List[List[float]]:
-        self.document_embeddings.embed(sentences)
-
-        text_embedding_list = [sentence.get_embedding().unsqueeze(0) for sentence in sentences]
-        text_embedding_tensor = torch.cat(text_embedding_list, 0).to(flair.device)
+        
+        if not self.attention:
+            self.document_embeddings.embed(sentences)
+            text_embedding_list = [sentence.get_embedding().unsqueeze(0) for sentence in sentences]
+            text_embedding_tensor = torch.cat(text_embedding_list, 0)
+            label_scores = self.decoder(text_embedding_tensor)
 
-        label_scores = self.decoder(text_embedding_tensor)
+        else:
+            lstm_outputs = self.document_embeddings.embed(sentences, return_sequences = True) 
+            text_embedding_tensor = self.attention_model.forward(attention_candidates = lstm_outputs, attention_size=lstm_outputs.size(-1)).squeeze(dim=1)
+            label_scores = self.decoder(text_embedding_tensor)
+            
 
         return label_scores
 
@@ -84,6 +136,7 @@ def save_checkpoint(self, model_file: Union[str, Path], optimizer_state: dict, s
             'state_dict': self.state_dict(),
             'document_embeddings': self.document_embeddings,
             'label_dictionary': self.label_dictionary,
+        
             'multi_label': self.multi_label,
             'optimizer_state_dict': optimizer_state,
             'scheduler_state_dict': scheduler_state,
@@ -266,4 +319,4 @@ def load(model: str):
             model_file = cached_path(base_path, cache_dir=cache_dir)
 
         if model_file is not None:
-            return TextClassifier.load_from_file(model_file)
+            return TextClassifier.load_from_file(model_file)
\ No newline at end of file
diff --git a/train.py b/train.py
deleted file mode 100644
--- a/train.py
+++ /dev/null
@@ -1,55 +0,0 @@
-from typing import List
-
-from flair.data_fetcher import NLPTaskDataFetcher, NLPTask
-from flair.data import TaggedCorpus
-from flair.embeddings import TokenEmbeddings, WordEmbeddings, StackedEmbeddings, CharLMEmbeddings, CharacterEmbeddings
-from flair.training_utils import EvaluationMetric
-from flair.visual.training_curves import Plotter
-
-# 1. get the corpus
-corpus: TaggedCorpus = NLPTaskDataFetcher.load_corpus(NLPTask.UD_ENGLISH)
-print(corpus)
-
-# 2. what tag do we want to predict?
-tag_type = 'pos'
-
-# 3. make the tag dictionary from the corpus
-tag_dictionary = corpus.make_tag_dictionary(tag_type=tag_type)
-print(tag_dictionary.idx2item)
-
-# initialize embeddings
-embedding_types: List[TokenEmbeddings] = [
-
-    WordEmbeddings('glove'),
-
-    # comment in this line to use character embeddings
-    # CharacterEmbeddings(),
-
-    # comment in these lines to use contextual string embeddings
-    #
-    # CharLMEmbeddings('news-forward'),
-    #
-    # CharLMEmbeddings('news-backward'),
-]
-
-embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embedding_types)
-
-# initialize sequence tagger
-from flair.models import SequenceTagger
-
-tagger: SequenceTagger = SequenceTagger(hidden_size=256,
-                                        embeddings=embeddings,
-                                        tag_dictionary=tag_dictionary,
-                                        tag_type=tag_type,
-                                        use_crf=True)
-
-# initialize trainer
-from flair.trainers import ModelTrainer
-trainer: ModelTrainer = ModelTrainer(tagger, corpus)
-
-trainer.train('resources/taggers/example-ner', EvaluationMetric.MICRO_F1_SCORE, learning_rate=0.1, mini_batch_size=32,
-              max_epochs=20, test_mode=True)
-
-plotter = Plotter()
-plotter.plot_training_curves('resources/taggers/example-ner/loss.tsv')
-plotter.plot_weights('resources/taggers/example-ner/weights.txt')
\ No newline at end of file


Test patch: