Repository: flairNLP/flair

Issue:
Problem with max_sequence_length in BertEmbeddings
Currently, `BertEmbeddings` does not account for the maximum sequence length supported by the underlying (`transformers`) `BertModel`. Since BERT creates subtokens, it becomes somewhat challenging to check sequence-length and trim sentence externally before feeding it to `BertEmbeddings` in `flair`.

I see a problem in https://github.com/flairNLP/flair/blob/master/flair/embeddings.py#L2678--L2687

```python
        # first, find longest sentence in batch
        longest_sentence_in_batch: int = len(
            max(
                [
                    self.tokenizer.tokenize(sentence.to_tokenized_string())
                    for sentence in sentences
                ],
                key=len,
            )
        )
```

This is passed to 
```python
        # prepare id maps for BERT model
        features = self._convert_sentences_to_features(
            sentences, longest_sentence_in_batch
        )
```
which sets `max_sequence_length` in: 

https://github.com/flairNLP/flair/blob/master/flair/embeddings.py#L2620-L2622
```python
    _convert_sentences_to_features(
        self, sentences, max_sequence_length: int
    )
````

But this does not account for or check the max-sequence-length supported by the BERT model, which is accessible in either of the above functions through ` self.model.config.max_position_embeddings`.


Discussion:
Hi @ayushjaiswal we are in the process of refactoring the transformer-based embeddings classes. See #1494. Instead of separate classes for each transformer embedding, we will have a unified class that gets the transformer model as string in the constructor. So initialization will be like this: 

```python
# example sentence
sentence = Sentence('The grass is green')

# a BERT model
embeddings = TransformerWordEmbeddings(model="bert-base-uncased", layers="-1", pooling_operation='first')
embeddings.embed(sentence)

# a roBERTa model
embeddings = TransformerWordEmbeddings(model="distilroberta-base", layers="-1", pooling_operation='first')
embeddings.embed(sentence)
```

There is now also a corresponding `TransformerDocumentEmbeddings` class in case you want document embeddings out of the transformer.

We're also looking at different ways for handling overlong sequences as part of the refactoring. We will add handling for this soon.
@alanakbik Thanks for the quick response! Great to hear about the refactoring and handling of overlong sequences. `self.model.config.max_position_embeddings` definitely needs to be accounted for so that the input sequence during forward pass of the `BertModel` does not encounter sequences of length greater than that. Currently, when the length does exceed the limit, a `RuntimeError` occurs caused by a `CUDA AssertionError` which corrupts the CUDA context and requires re-initialization of the CUDA session. Even if the input sequence is trimmed, I guess it will create a problem with assigning embeddings to `Sentence` tokens. It seems somewhat tricky 😅 
@alanakbik 
Maybe a sliding window approach, as implemented [here](https://github.com/ThilinaRajapakse/simpletransformers/blob/e289efa09e280955207fefbec8aedcced8449be3/simpletransformers/language_modeling/language_modeling_utils.py) , might be a good idea to tackle the length limitation of BERT.
I've resorted a lot to using the linked package instead of flair, solely for this feature, as the results seem to be better compared to simply truncating the sentences.

Would love to see this feature in flair!

Thanks for the pointer - yes this looks promising so we might integrate it!
Looking forward to this 😄 
@alanakbik is there any update on this? 🙂 

Resolving patch:
diff --git a/flair/embeddings/token.py b/flair/embeddings/token.py
--- a/flair/embeddings/token.py
+++ b/flair/embeddings/token.py
@@ -812,7 +812,8 @@ def __init__(
         pooling_operation: str = "first",
         batch_size: int = 1,
         use_scalar_mix: bool = False,
-        fine_tune: bool = False
+        fine_tune: bool = False,
+        allow_long_sentences: bool = False
     ):
         """
         Bidirectional transformer embeddings of words from various transformer architectures.
@@ -833,6 +834,15 @@ def __init__(
         config = AutoConfig.from_pretrained(model, output_hidden_states=True)
         self.model = AutoModel.from_pretrained(model, config=config)
 
+        self.allow_long_sentences = allow_long_sentences
+
+        if allow_long_sentences:
+            self.max_subtokens_sequence_length = self.tokenizer.model_max_length
+            self.stride = self.tokenizer.model_max_length//2
+        else:
+            self.max_subtokens_sequence_length = None
+            self.stride = 0
+
         # model name
         self.name = 'transformer-word-' + str(model)
 
@@ -893,8 +903,7 @@ def _remove_special_markup(text: str):
         return text
 
     def _get_processed_token_text(self, token: Token) -> str:
-        pieces = self.tokenizer.convert_ids_to_tokens(
-            self.tokenizer.encode(token.text, add_special_tokens=False))
+        pieces = self.tokenizer.tokenize(token.text)
         token_text = ''
         for piece in pieces:
             token_text += self._remove_special_markup(piece)
@@ -908,87 +917,62 @@ def _add_embeddings_to_sentences(self, sentences: List[Sentence]):
         subtokenized_sentences = []
         subtokenized_sentences_token_lengths = []
 
-        for sentence in sentences:
+        sentence_parts_lengths = []
+
+        # TODO: keep for backwards compatibility, but remove in future
+        # some pretrained models do not have this property, applying default settings now.
+        # can be set manually after loading the model.
+        if not hasattr(self, 'max_subtokens_sequence_length'):
+            self.max_subtokens_sequence_length = None
+            self.allow_long_sentences = False
+            self.stride = 0
 
+        for sentence in sentences:
             tokenized_string = sentence.to_tokenized_string()
 
             # method 1: subtokenize sentence
             # subtokenized_sentence = self.tokenizer.encode(tokenized_string, add_special_tokens=True)
 
             # method 2:
-            ids = self.tokenizer.encode(tokenized_string, add_special_tokens=False)
-            subtokenized_sentence = self.tokenizer.build_inputs_with_special_tokens(ids)
-
-            subtokenized_sentences.append(torch.tensor(subtokenized_sentence, dtype=torch.long))
-            subtokens = self.tokenizer.convert_ids_to_tokens(subtokenized_sentence)
-
-            word_iterator = iter(sentence)
-            token = next(word_iterator)
-            token_text = self._get_processed_token_text(token)
-
-            token_subtoken_lengths = []
-            reconstructed_token = ''
-            subtoken_count = 0
-
-            # iterate over subtokens and reconstruct tokens
-            for subtoken_id, subtoken in enumerate(subtokens):
-
-                # remove special markup
-                subtoken = self._remove_special_markup(subtoken)
-
-                # check if subtoken is special begin token ([CLS] or similar)
-                if subtoken in self.special_tokens and subtoken_id == 0:
-                    continue
-
-                # some BERT tokenizers somehow omit words - in such cases skip to next token
-                if subtoken_count == 0 and not token_text.startswith(subtoken.lower()):
-                    token_subtoken_lengths.append(0)
-                    token = next(word_iterator)
-                    token_text = self._get_processed_token_text(token)
-
-                subtoken_count += 1
+            # transformer specific tokenization
+            subtokenized_sentence = self.tokenizer.tokenize(tokenized_string)
+            token_subtoken_lengths = self.reconstruct_tokens_from_subtokens(sentence, subtokenized_sentence)
+            subtokenized_sentences_token_lengths.append(token_subtoken_lengths)
 
-                # append subtoken to reconstruct token
-                reconstructed_token = reconstructed_token + subtoken
+            subtoken_ids_sentence = self.tokenizer.convert_tokens_to_ids(subtokenized_sentence)
 
-                # check if reconstructed token is the same as current token
-                if reconstructed_token.lower() == token_text:
+            nr_sentence_parts = 0
 
-                    # if so, add subtoken count
-                    token_subtoken_lengths.append(subtoken_count)
+            while subtoken_ids_sentence:
+                nr_sentence_parts += 1
+                encoded_inputs = self.tokenizer.prepare_for_model(subtoken_ids_sentence,
+                                                                  max_length=self.max_subtokens_sequence_length,
+                                                                  stride=self.stride,
+                                                                  return_overflowing_tokens=self.allow_long_sentences)
 
-                    # reset subtoken count and reconstructed token
-                    reconstructed_token = ''
-                    subtoken_count = 0
+                subtoken_ids_split_sentence = encoded_inputs['input_ids']
+                subtokenized_sentences.append(torch.tensor(subtoken_ids_split_sentence, dtype=torch.long))
 
-                    # get next token
-                    if len(token_subtoken_lengths) < len(sentence):
-                        token = next(word_iterator)
-                        token_text = self._get_processed_token_text(token)
-                    # break from loop if all tokens are accounted for
-                    else:
-                        break
 
-            # check if all tokens were matched to subtokens
-            if token != sentence[-1]:
-                log.error(f"Tokenization MISMATCH in sentence '{sentence.to_tokenized_string()}'")
-                log.error(f"Last matched: '{token}'")
-                log.error(f"Last sentence: '{sentence[-1]}'")
-                log.error(f"subtokenized: '{subtokens}'")
+                if 'overflowing_tokens' in encoded_inputs:
+                    subtoken_ids_sentence = encoded_inputs['overflowing_tokens']
+                else:
+                    subtoken_ids_sentence = None
 
-            subtokenized_sentences_token_lengths.append(token_subtoken_lengths)
+            sentence_parts_lengths.append(nr_sentence_parts)
 
         # find longest sentence in batch
         longest_sequence_in_batch: int = len(max(subtokenized_sentences, key=len))
 
+        total_sentence_parts = sum(sentence_parts_lengths)
         # initialize batch tensors and mask
         input_ids = torch.zeros(
-            [len(sentences), longest_sequence_in_batch],
+            [total_sentence_parts, longest_sequence_in_batch],
             dtype=torch.long,
             device=flair.device,
         )
         mask = torch.zeros(
-            [len(sentences), longest_sequence_in_batch],
+            [total_sentence_parts, longest_sequence_in_batch],
             dtype=torch.long,
             device=flair.device,
         )
@@ -999,6 +983,10 @@ def _add_embeddings_to_sentences(self, sentences: List[Sentence]):
 
         # put encoded batch through transformer model to get all hidden states of all encoder layers
         hidden_states = self.model(input_ids, attention_mask=mask)[-1]
+        # make the tuple a tensor; makes working with it easier.
+        hidden_states = torch.stack(hidden_states)
+
+        sentence_idx_offset = 0
 
         # gradients are enabled if fine-tuning is enabled
         gradient_context = torch.enable_grad() if (self.fine_tune and self.training) else torch.no_grad()
@@ -1006,7 +994,18 @@ def _add_embeddings_to_sentences(self, sentences: List[Sentence]):
         with gradient_context:
 
             # iterate over all subtokenized sentences
-            for sentence_idx, (sentence, subtoken_lengths) in enumerate(zip(sentences, subtokenized_sentences_token_lengths)):
+            for sentence_idx, (sentence, subtoken_lengths, nr_sentence_parts) in enumerate(zip(sentences, subtokenized_sentences_token_lengths, sentence_parts_lengths)):
+
+                sentence_hidden_state = hidden_states[:, sentence_idx + sentence_idx_offset, ...]
+
+                for i in range(1, nr_sentence_parts):
+                    sentence_idx_offset += 1
+                    remainder_sentence_hidden_state = hidden_states[:, sentence_idx + sentence_idx_offset, ...]
+                    # remove stride_size//2 at end of sentence_hidden_state, and half at beginning of remainder,
+                    # in order to get some context into the embeddings of these words.
+                    # also don't include the embedding of the extra [CLS] and [SEP] tokens.
+                    sentence_hidden_state = torch.cat((sentence_hidden_state[:, :-1-self.stride//2, :],
+                                                       remainder_sentence_hidden_state[:, 1 + self.stride//2:, :]), 1)
 
                 subword_start_idx = self.begin_offset
 
@@ -1024,7 +1023,7 @@ def _add_embeddings_to_sentences(self, sentences: List[Sentence]):
 
                     # get states from all selected layers, aggregate with pooling operation
                     for layer in self.layer_indexes:
-                        current_embeddings = hidden_states[layer][sentence_idx][subword_start_idx:subword_end_idx]
+                        current_embeddings = sentence_hidden_state[layer][subword_start_idx:subword_end_idx]
 
                         if self.pooling_operation == "first":
                             final_embedding: torch.FloatTensor = current_embeddings[0]
@@ -1055,6 +1054,59 @@ def _add_embeddings_to_sentences(self, sentences: List[Sentence]):
 
                     subword_start_idx += number_of_subtokens
 
+    def reconstruct_tokens_from_subtokens(self, sentence, subtokens):
+        word_iterator = iter(sentence)
+        token = next(word_iterator)
+        token_text = self._get_processed_token_text(token)
+        token_subtoken_lengths = []
+        reconstructed_token = ''
+        subtoken_count = 0
+        # iterate over subtokens and reconstruct tokens
+        for subtoken_id, subtoken in enumerate(subtokens):
+
+            # remove special markup
+            subtoken = self._remove_special_markup(subtoken)
+
+            # TODO check if this is necessary is this method is called before prepare_for_model
+            # check if reconstructed token is special begin token ([CLS] or similar)
+            if subtoken in self.special_tokens and subtoken_id == 0:
+                continue
+
+            # some BERT tokenizers somehow omit words - in such cases skip to next token
+            if subtoken_count == 0 and not token_text.startswith(subtoken.lower()):
+                token_subtoken_lengths.append(0)
+                token = next(word_iterator)
+                token_text = self._get_processed_token_text(token)
+
+            subtoken_count += 1
+
+            # append subtoken to reconstruct token
+            reconstructed_token = reconstructed_token + subtoken
+
+            # check if reconstructed token is the same as current token
+            if reconstructed_token.lower() == token_text:
+
+                # if so, add subtoken count
+                token_subtoken_lengths.append(subtoken_count)
+
+                # reset subtoken count and reconstructed token
+                reconstructed_token = ''
+                subtoken_count = 0
+
+                # break from loop if all tokens are accounted for
+                if len(token_subtoken_lengths) < len(sentence):
+                    token = next(word_iterator)
+                    token_text = self._get_processed_token_text(token)
+                else:
+                    break
+        # check if all tokens were matched to subtokens
+        if token != sentence[-1]:
+            log.error(f"Tokenization MISMATCH in sentence '{sentence.to_tokenized_string()}'")
+            log.error(f"Last matched: '{token}'")
+            log.error(f"Last sentence: '{sentence[-1]}'")
+            log.error(f"subtokenized: '{subtokens}'")
+        return token_subtoken_lengths
+
     def train(self, mode=True):
         # if fine-tuning is not enabled (i.e. a "feature-based approach" used), this
         # module should never be in training mode


Test patch: