Repository: RasaHQ/rasa
Pull request: Better handling of sequences longer than maximum sequence length for BERT style models
**Proposed changes**:
- Closes https://github.com/RasaHQ/rasa/issues/6222

**Status (please check what you already did)**:
- [x] added some tests for the functionality
- [ ] updated the documentation
- [x] updated the changelog (please check [changelog](https://github.com/RasaHQ/rasa/tree/master/changelog) for instructions)
- [x] reformat files using `black` (please check [Readme](https://github.com/RasaHQ/rasa#code-style) for instructions)


Problem statement:
Problem with integrating BERT (tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0,512] = 512 is not in [0, 512))
**Rasa version**: 1.10.7 (also tried with 1.8.3)

**Python version**: 3.7

**Operating system** (windows, osx, ...): Linux and Windows

**Issue**:
I have a problem with integrating BERT into the configuration of my model. Namely, at the very beginning of training it gives the error message shown below. There is also a config.yml file I used. 

It is interesting that I am not getting this error when using a very small train dataset, so the model is able to finish with training successfully.

**Error (including full traceback)**:
```
2020-07-16 13:05:16 INFO     rasa.nlu.model  - Starting to train component HFTransformersNLP
Traceback (most recent call last):
  File "/home/ubuntu/anaconda3/bin/rasa", line 8, in <module>
    sys.exit(main())
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/__main__.py", line 92, in main
    cmdline_arguments.func(cmdline_arguments)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/cli/train.py", line 140, in train_nlu
    persist_nlu_training_data=args.persist_nlu_data,
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/train.py", line 414, in train_nlu
    persist_nlu_training_data,
  File "uvloop/loop.pyx", line 1456, in uvloop.loop.Loop.run_until_complete
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/train.py", line 453, in _train_nlu_async
    persist_nlu_training_data=persist_nlu_training_data,
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/train.py", line 482, in _train_nlu_with_validated_data
    persist_nlu_training_data=persist_nlu_training_data,
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/train.py", line 90, in train
    interpreter = trainer.train(training_data, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/model.py", line 191, in train
    updates = component.train(working_data, self.config, **context)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 472, in train
    batch_docs = self._get_docs_for_batch(batch_messages, attribute)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 423, in _get_docs_for_batch
    ) = self._get_model_features_for_batch(batch_token_ids)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 387, in _get_model_features_for_batch
    batch_attention_mask, padded_token_ids
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/rasa/nlu/utils/hugging_face/hf_transformers.py", line 347, in _compute_batch_sequence_features
    np.array(padded_token_ids), attention_mask=np.array(batch_attention_mask)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 708, in call
    outputs = self.bert(inputs, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 570, in call
    embedding_output = self.embeddings([input_ids, position_ids, token_type_ids, inputs_embeds], training=training)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 150, in call
    return self._embedding(inputs, training=training)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/transformers/modeling_tf_bert.py", line 173, in _embedding
    position_embeddings = self.position_embeddings(position_ids)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 822, in __call__
    outputs = self.call(cast_inputs, *args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/embeddings.py", line 184, in call
    out = embedding_ops.embedding_lookup(self.embeddings, inputs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/embedding_ops.py", line 323, in embedding_lookup
    transform_fn=None)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/embedding_ops.py", line 137, in _embedding_lookup_and_transform
    array_ops.gather(params[0], ids, name=name), ids, max_norm)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/util/dispatch.py", line 180, in wrapper
    return target(*args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py", line 4106, in gather
    return params.sparse_read(indices, name=name)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py", line 645, in sparse_read
    self._handle, indices, dtype=self._dtype, name=name)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_resource_variable_ops.py", line 556, in resource_gather
    _ops.raise_from_not_ok_status(e, name)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 6606, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0,512] = 512 is not in [0, 512) [Op:ResourceGather] name: tf_bert_model/bert/embeddings/position_embeddings/embedding_lookup/
```

**Command or request that led to error**:
```
rasa train nlu
```

**Content of configuration file (config.yml)** (if relevant):
```yml

language: en
num_threads: 8
pipeline:
- name: HFTransformersNLP
  model_weights: "bert-base-uncased"
  model_name: "bert"
- name: LanguageModelTokenizer
- name: LanguageModelFeaturizer
- name: "CRFEntityExtractor"
- name: "EmbeddingIntentClassifier"

```

Hints from issue discussion:
Thanks for raising this issue, @tabergma will get back to you about it soon✨

###### Please also check out the [docs](https://rasa.com/docs/) and the [forum](https://forum.rasa.com/) in case your issue was raised there too 🤗
@MajaRolevski Thanks for raising the issue. How small is your "very small dataset"? And how does your "normal" dataset look like that resolves in this issue?
@tabergma Thank you for the answer. My "very small dataset" consists of 19 classes (intents); each class has 10-30 sentences, except one that has 347 sentences and the other with 1897. On the other hand, the "normal" dataset contains 19 classes, where each class has on average 70-100 sentences, except one with 1565 and the other with 41564 sentences.
@MajaRolevski Looks like your dataset has some sentences which are very long(>512 tokens). We'll add a fix soon so that the training doesn't break because of this but I would also suggest you to sanitize your data and possibly clean up such sentences because they are unusual and could be very noisy for training.
We are facing the same issue. Our response selector has some rather long answers, so we run into the BERT token limit as well. When I shorten my answers, the model trains fine again. It is hard to find a good cut off length for the answers though with the different tokenizer and libs involved.

I am looking forward to fix for this :-)

Commit message: more tests

Patch to review:
diff --git a/changelog/6222.bugfix.md b/changelog/6222.bugfix.md
new file mode 100644
index 000000000000..1ba2df6e9dc7
--- /dev/null
+++ b/changelog/6222.bugfix.md
@@ -0,0 +1,5 @@
+Better handling of input sequences longer than the maximum sequence length that the `HFTransformersNLP` models can handle.
+
+During training, messages with longer sequence length should result in an error, whereas during inference they are
+gracefully handled but a debug message is logged. Ideally, passing messages longer than the acceptable maximum sequence 
+lengths of each model should be avoided.
\ No newline at end of file
diff --git a/rasa/nlu/constants.py b/rasa/nlu/constants.py
index 430ecf86eb3f..b71ac254fed4 100644
--- a/rasa/nlu/constants.py
+++ b/rasa/nlu/constants.py
@@ -74,3 +74,5 @@
 VALID_FEATURE_TYPES = [FEATURE_TYPE_SEQUENCE, FEATURE_TYPE_SENTENCE]
 
 FEATURIZER_CLASS_ALIAS = "alias"
+
+NO_LENGTH_RESTRICTION = -1
diff --git a/rasa/nlu/utils/hugging_face/hf_transformers.py b/rasa/nlu/utils/hugging_face/hf_transformers.py
index 0907a004a8be..fbaea861298a 100644
--- a/rasa/nlu/utils/hugging_face/hf_transformers.py
+++ b/rasa/nlu/utils/hugging_face/hf_transformers.py
@@ -20,8 +20,18 @@
     SENTENCE_FEATURES,
     SEQUENCE_FEATURES,
     NUMBER_OF_SUB_TOKENS,
+    NO_LENGTH_RESTRICTION,
 )
 
+MAX_SEQUENCE_LENGTHS = {
+    "bert": 512,
+    "gpt": 512,
+    "gpt2": 512,
+    "xlnet": NO_LENGTH_RESTRICTION,
+    "distilbert": 512,
+    "roberta": 512,
+}
+
 logger = logging.getLogger(__name__)
 
 
@@ -44,19 +54,22 @@ class HFTransformersNLP(Component):
         "cache_dir": None,
     }
 
-    def __init__(self, component_config: Optional[Dict[Text, Any]] = None) -> None:
+    def __init__(
+        self,
+        component_config: Optional[Dict[Text, Any]] = None,
+        skip_model_load: bool = False,
+    ) -> None:
         super(HFTransformersNLP, self).__init__(component_config)
 
-        self._load_model()
+        self._load_model_metadata()
+        self._load_model_instance(skip_model_load)
         self.whitespace_tokenizer = WhitespaceTokenizer()
 
-    def _load_model(self) -> None:
-        """Try loading the model"""
+    def _load_model_metadata(self) -> None:
 
         from rasa.nlu.utils.hugging_face.registry import (
             model_class_dict,
             model_weights_defaults,
-            model_tokenizer_dict,
         )
 
         self.model_name = self.component_config["model_name"]
@@ -64,7 +77,7 @@ def _load_model(self) -> None:
         if self.model_name not in model_class_dict:
             raise KeyError(
                 f"'{self.model_name}' not a valid model name. Choose from "
-                f"{str(list(model_class_dict.keys()))}or create"
+                f"{str(list(model_class_dict.keys()))} or create"
                 f"a new class inheriting from this class to support your model."
             )
 
@@ -78,6 +91,24 @@ def _load_model(self) -> None:
             )
             self.model_weights = model_weights_defaults[self.model_name]
 
+        self.max_model_sequence_length = MAX_SEQUENCE_LENGTHS[self.model_name]
+
+    def _load_model_instance(self, skip_model_load: bool) -> None:
+        """Try loading the model instance
+
+        Args:
+            skip_model_load: Skip loading the model instances to save time. This should be True only for pytests
+        """
+
+        if skip_model_load:
+            # This should be True only during pytests
+            return
+
+        from rasa.nlu.utils.hugging_face.registry import (
+            model_class_dict,
+            model_tokenizer_dict,
+        )
+
         logger.debug(f"Loading Tokenizer and Model for {self.model_name}")
 
         self.tokenizer = model_tokenizer_dict[self.model_name].from_pretrained(
@@ -109,8 +140,7 @@ def required_packages(cls) -> List[Text]:
         return ["transformers"]
 
     def _lm_tokenize(self, text: Text) -> Tuple[List[int], List[Text]]:
-        """
-        Pass the text through the tokenizer of the language model.
+        """Pass the text through the tokenizer of the language model.
 
         Args:
             text: Text to be tokenized.
@@ -128,8 +158,7 @@ def _lm_tokenize(self, text: Text) -> Tuple[List[int], List[Text]]:
     def _add_lm_specific_special_tokens(
         self, token_ids: List[List[int]]
     ) -> List[List[int]]:
-        """Add language model specific special tokens which were used during their
-        training.
+        """Add language model specific special tokens which were used during their training.
 
         Args:
             token_ids: List of token ids for each example in the batch.
@@ -276,64 +305,96 @@ def _get_token_ids_for_batch(
         return batch_tokens, batch_token_ids
 
     @staticmethod
-    def _compute_attention_mask(actual_sequence_lengths: List[int]) -> np.ndarray:
+    def _compute_attention_mask(
+        actual_sequence_lengths: List[int], max_input_sequence_length: int
+    ) -> np.ndarray:
         """Compute a mask for padding tokens.
 
         This mask will be used by the language model so that it does not attend to
         padding tokens.
 
         Args:
-            actual_sequence_lengths: List of length of each example without any padding
+            actual_sequence_lengths: List of length of each example without any padding.
+            max_input_sequence_length: Maximum length of a sequence that will be present in the input batch. This is
+            after taking into consideration the maximum input sequence the model can handle. Hence it can never be
+            greater than self.max_model_sequence_length in case the model applies length restriction.
 
         Returns:
             Computed attention mask, 0 for padding and 1 for non-padding tokens.
         """
 
         attention_mask = []
-        max_seq_length = max(actual_sequence_lengths)
+
         for actual_sequence_length in actual_sequence_lengths:
             # add 1s for present tokens, fill up the remaining space up to max
             # sequence length with 0s (non-existing tokens)
-            padded_sequence = [1] * actual_sequence_length + [0] * (
-                max_seq_length - actual_sequence_length
+            padded_sequence = [1] * min(
+                actual_sequence_length, max_input_sequence_length
+            ) + [0] * (
+                max_input_sequence_length
+                - min(actual_sequence_length, max_input_sequence_length)
             )
             attention_mask.append(padded_sequence)
 
         attention_mask = np.array(attention_mask).astype(np.float32)
-
         return attention_mask
 
-    def _add_padding_to_batch(
+    def _extract_sequence_lengths(
         self, batch_token_ids: List[List[int]]
-    ) -> Tuple[List[int], List[List[int]]]:
+    ) -> Tuple[List[int], int]:
+
+        # Compute max length across examples
+        max_input_sequence_length = 0
+        actual_sequence_lengths = []
+
+        for example_token_ids in batch_token_ids:
+            sequence_length = len(example_token_ids)
+            actual_sequence_lengths.append(sequence_length)
+            max_input_sequence_length = max(
+                max_input_sequence_length, len(example_token_ids)
+            )
+
+        # Take into account the maximum sequence length the model can handle
+        max_input_sequence_length = (
+            max_input_sequence_length
+            if self.max_model_sequence_length == NO_LENGTH_RESTRICTION
+            else min(max_input_sequence_length, self.max_model_sequence_length)
+        )
+
+        return actual_sequence_lengths, max_input_sequence_length
+
+    def _add_padding_to_batch(
+        self, batch_token_ids: List[List[int]], max_sequence_length_model: int
+    ) -> List[List[int]]:
         """Add padding so that all examples in the batch are of the same length.
 
         Args:
             batch_token_ids: Batch of examples where each example is a non-padded list
             of token ids.
+            max_sequence_length_model: Maximum length of any input sequence in the batch
+            to be fed to the model.
 
         Returns:
             Padded batch with all examples of the same length.
         """
         padded_token_ids = []
-        # Compute max length across examples
-        max_seq_len = 0
-        actual_sequence_lengths = []
-
-        for example_token_ids in batch_token_ids:
-            actual_sequence_lengths.append(len(example_token_ids))
-            max_seq_len = max(max_seq_len, len(example_token_ids))
 
-        # Add padding according to max_seq_len
+        # Add padding according to max_sequence_length
         # Some models don't contain pad token, we use unknown token as padding token.
         # This doesn't affect the computation since we compute an attention mask
         # anyways.
         for example_token_ids in batch_token_ids:
+
+            # Truncate any longer sequences so that they can be fed to the model
+            if len(example_token_ids) > max_sequence_length_model:
+                example_token_ids = example_token_ids[:max_sequence_length_model]
+
             padded_token_ids.append(
                 example_token_ids
-                + [self.pad_token_id] * (max_seq_len - len(example_token_ids))
+                + [self.pad_token_id]
+                * (max_sequence_length_model - len(example_token_ids))
             )
-        return actual_sequence_lengths, padded_token_ids
+        return padded_token_ids
 
     @staticmethod
     def _extract_nonpadded_embeddings(
@@ -380,8 +441,95 @@ def _compute_batch_sequence_features(
         sequence_hidden_states = sequence_hidden_states.numpy()
         return sequence_hidden_states
 
+    def _validate_sequence_lengths(
+        self,
+        actual_sequence_lengths: List[int],
+        batch_examples: List[Message],
+        attribute: Text,
+        inference_mode: bool = False,
+    ) -> None:
+        """Validate if sequence lengths of all inputs are less the max sequence length the model can handle
+
+        This method should throw an error during training, whereas log a debug message during inference if
+        any of the input examples have a length greater than maximum sequence length allowed.
+
+        Args:
+            actual_sequence_lengths: original sequence length of all inputs
+            batch_examples: all message instances in the batch
+            attribute: attribute of message object to be processed
+            inference_mode: Whether this is during training or during inferencing
+        """
+        if self.max_model_sequence_length == NO_LENGTH_RESTRICTION:
+            # There is no restriction on sequence length from the model
+            return
+
+        for sequence_length, example in zip(actual_sequence_lengths, batch_examples):
+            if sequence_length > self.max_model_sequence_length:
+                if not inference_mode:
+                    raise RuntimeError(
+                        f"The sequence length of '{example.get(attribute)[:20]}...' "
+                        f"is too long({sequence_length} tokens) for the "
+                        f"model chosen {self.model_name} which has a maximum "
+                        f"sequence length of {self.max_model_sequence_length} tokens. Either "
+                        f"shorten the message or use a model which has no "
+                        f"restriction on input sequence length like XLNet."
+                    )
+                else:
+                    logger.debug(
+                        f"The sequence length of '{example.get(attribute)[:20]}...' "
+                        f"is too long({sequence_length} tokens) for the "
+                        f"model chosen {self.model_name} which has a maximum "
+                        f"sequence length of {self.max_model_sequence_length} tokens. "
+                        f"Downstream model predictions may be affected because of this."
+                    )
+
+    def _add_extra_padding(
+        self, sequence_embeddings: np.ndarray, actual_sequence_lengths: List[int]
+    ) -> np.ndarray:
+        """
+        Add extra zero padding to match the original sequence length.
+
+        This is only done if the input was truncated during the batch preparation of input for the model.
+        Args:
+            sequence_embeddings: Embeddings returned from the model
+            actual_sequence_lengths: original sequence length of all inputs
+
+        Returns:
+            Modified sequence embeddings with padding if necessary
+        """
+
+        if self.max_model_sequence_length == NO_LENGTH_RESTRICTION:
+            # No extra padding needed because there wouldn't have been any truncation in the first place
+            return sequence_embeddings
+
+        reshaped_sequence_embeddings = []
+        for index, embedding in enumerate(sequence_embeddings):
+            embedding_size = embedding.shape[-1]
+            if actual_sequence_lengths[index] > self.max_model_sequence_length:
+                embedding = np.concatenate(
+                    [
+                        embedding,
+                        np.zeros(
+                            (
+                                actual_sequence_lengths[index]
+                                - self.max_model_sequence_length,
+                                embedding_size,
+                            ),
+                            dtype=np.float32,
+                        ),
+                    ]
+                )
+            reshaped_sequence_embeddings.append(embedding)
+
+        return np.array(reshaped_sequence_embeddings)
+
     def _get_model_features_for_batch(
-        self, batch_token_ids: List[List[int]], batch_tokens: List[List[Token]]
+        self,
+        batch_token_ids: List[List[int]],
+        batch_tokens: List[List[Token]],
+        batch_examples: List[Message],
+        attribute: Text,
+        inference_mode: bool = False,
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Compute dense features of each example in the batch.
 
@@ -394,6 +542,10 @@ def _get_model_features_for_batch(
 
         Args:
             batch_token_ids: List of token ids of each example in the batch.
+            batch_tokens: List of token objects for each example in the batch.
+            batch_examples: List of examples in the batch.
+            attribute: attribute of the Message object to be processed.
+            inference_mode: Whether the call is during training or during inference.
 
         Returns:
             Sentence and token level dense representations.
@@ -403,13 +555,27 @@ def _get_model_features_for_batch(
             batch_token_ids
         )
 
-        # Let's first add padding so that whole batch can be fed to the model
-        actual_sequence_lengths, padded_token_ids = self._add_padding_to_batch(
-            batch_token_ids_augmented
+        # Compute sequence lengths for all examples
+        (
+            actual_sequence_lengths,
+            max_input_sequence_length,
+        ) = self._extract_sequence_lengths(batch_token_ids_augmented)
+
+        # Validate that all sequences can be processed based on their sequence lengths and
+        # the maximum sequence length the model can handle
+        self._validate_sequence_lengths(
+            actual_sequence_lengths, batch_examples, attribute, inference_mode
+        )
+
+        # Add padding so that whole batch can be fed to the model
+        padded_token_ids = self._add_padding_to_batch(
+            batch_token_ids_augmented, max_input_sequence_length
         )
 
         # Compute attention mask based on actual_sequence_length
-        batch_attention_mask = self._compute_attention_mask(actual_sequence_lengths)
+        batch_attention_mask = self._compute_attention_mask(
+            actual_sequence_lengths, max_input_sequence_length
+        )
 
         # Get token level features from the model
         sequence_hidden_states = self._compute_batch_sequence_features(
@@ -427,6 +593,12 @@ def _get_model_features_for_batch(
             sequence_embeddings,
         ) = self._post_process_sequence_embeddings(sequence_nonpadded_embeddings)
 
+        # Pad zeros for examples which were truncated in inference mode.
+        # This is intentionally done after sentence embeddings have been extracted so that they are not affected
+        sequence_embeddings = self._add_extra_padding(
+            sequence_embeddings, actual_sequence_lengths
+        )
+
         # shape of matrix for all sequence embeddings
         batch_dim = len(sequence_embeddings)
         seq_dim = max(e.shape[0] for e in sequence_embeddings)
@@ -449,7 +621,10 @@ def _get_model_features_for_batch(
         return sentence_embeddings, sequence_final_embeddings
 
     def _get_docs_for_batch(
-        self, batch_examples: List[Message], attribute: Text
+        self,
+        batch_examples: List[Message],
+        attribute: Text,
+        inference_mode: bool = False,
     ) -> List[Dict[Text, Any]]:
         """Compute language model docs for all examples in the batch.
 
@@ -458,6 +633,8 @@ def _get_docs_for_batch(
             need to be computed.
             attribute: Property of message to be processed, one of ``TEXT`` or
             ``RESPONSE``.
+            inference_mode: Whether the call is during inference or during training.
+
 
         Returns:
             List of language model docs for each message in batch.
@@ -470,7 +647,9 @@ def _get_docs_for_batch(
         (
             batch_sentence_features,
             batch_sequence_features,
-        ) = self._get_model_features_for_batch(batch_token_ids, batch_tokens)
+        ) = self._get_model_features_for_batch(
+            batch_token_ids, batch_tokens, batch_examples, attribute, inference_mode
+        )
 
         # A doc consists of
         # {'token_ids': ..., 'tokens': ..., 'sequence_features': ...,
@@ -538,5 +717,5 @@ def process(self, message: Message, **kwargs: Any) -> None:
 
         message.set(
             LANGUAGE_MODEL_DOCS[TEXT],
-            self._get_docs_for_batch([message], attribute=TEXT)[0],
+            self._get_docs_for_batch([message], attribute=TEXT, inference_mode=True)[0],
         )
diff --git a/tests/nlu/utils/test_hf_transformers.py b/tests/nlu/utils/test_hf_transformers.py
new file mode 100644
index 000000000000..3edcb6e22f0c
--- /dev/null
+++ b/tests/nlu/utils/test_hf_transformers.py
@@ -0,0 +1,126 @@
+import pytest
+import numpy as np
+from typing import List, Text
+import logging
+
+from rasa.nlu.utils.hugging_face.hf_transformers import HFTransformersNLP
+from rasa.nlu.training_data import Message
+
+
+@pytest.mark.parametrize(
+    "input_sequence_length, model_name, should_overflow",
+    [(20, "bert", False), (1000, "bert", True), (1000, "xlnet", False)],
+)
+def test_sequence_length_overflow_train(
+    input_sequence_length: int, model_name: Text, should_overflow: bool
+):
+
+    component = HFTransformersNLP({"model_name": model_name}, skip_model_load=True)
+    message = Message.build(text=" ".join(["hi"] * input_sequence_length))
+    if should_overflow:
+        with pytest.raises(RuntimeError):
+            component._validate_sequence_lengths(
+                [input_sequence_length], [message], "text", inference_mode=False
+            )
+    else:
+        component._validate_sequence_lengths(
+            [input_sequence_length], [message], "text", inference_mode=False
+        )
+
+
+@pytest.mark.parametrize(
+    "sequence_embeddings, actual_sequence_lengths, model_name, padding_needed",
+    [
+        (np.ones((1, 512, 5)), [1000], "bert", True),
+        (np.ones((1, 512, 5)), [1000], "xlnet", False),
+        (np.ones((1, 256, 5)), [256], "bert", False),
+    ],
+)
+def test_long_sequences_extra_padding(
+    sequence_embeddings: np.ndarray,
+    actual_sequence_lengths: List[int],
+    model_name: Text,
+    padding_needed: bool,
+):
+
+    component = HFTransformersNLP({"model_name": model_name}, skip_model_load=True)
+    modified_sequence_embeddings = component._add_extra_padding(
+        sequence_embeddings, actual_sequence_lengths
+    )
+    if not padding_needed:
+        assert np.all(modified_sequence_embeddings) == np.all(sequence_embeddings)
+    else:
+        assert modified_sequence_embeddings.shape[1] == actual_sequence_lengths[0]
+        assert (
+            modified_sequence_embeddings[0].shape[-1]
+            == sequence_embeddings[0].shape[-1]
+        )
+        zero_embeddings = modified_sequence_embeddings[0][
+            sequence_embeddings.shape[1] :
+        ]
+        assert np.all(zero_embeddings == 0)
+
+
+@pytest.mark.parametrize(
+    "token_ids, max_sequence_length_model, resulting_length, padding_added",
+    [
+        ([[1] * 200], 512, 512, True),
+        ([[1] * 700], 512, 512, False),
+        ([[1] * 200], 200, 200, False),
+    ],
+)
+def test_input_padding(
+    token_ids: List[List[int]],
+    max_sequence_length_model: int,
+    resulting_length: int,
+    padding_added: bool,
+):
+    component = HFTransformersNLP({"model_name": "bert"}, skip_model_load=True)
+    component.pad_token_id = 0
+    padded_input = component._add_padding_to_batch(token_ids, max_sequence_length_model)
+    assert len(padded_input[0]) == resulting_length
+    if padding_added:
+        original_length = len(token_ids[0])
+        assert np.all(np.array(padded_input[0][original_length:]) == 0)
+
+
+@pytest.mark.parametrize(
+    "sequence_length, model_name, should_overflow",
+    [(1000, "bert", True), (256, "bert", False)],
+)
+@pytest.mark.skip_on_windows
+def test_log_longer_sequence(
+    sequence_length: int, model_name: Text, should_overflow: bool, caplog
+):
+    transformers_config = {"model_name": model_name}
+
+    transformers_nlp = HFTransformersNLP(transformers_config)
+
+    text = " ".join(["hi"] * sequence_length)
+    message = Message(text)
+
+    caplog.set_level(logging.DEBUG)
+    transformers_nlp.process(message)
+    if should_overflow:
+        print(caplog.text)
+        assert "hi hi hi" in caplog.text
+    assert message.get("text_language_model_doc") is not None
+
+
+@pytest.mark.parametrize(
+    "actual_sequence_length, max_input_sequence_length, zero_start_index",
+    [(256, 512, 256), (700, 700, 700), (700, 512, 512)],
+)
+def test_attention_mask(
+    actual_sequence_length: int, max_input_sequence_length: int, zero_start_index: int
+):
+    component = HFTransformersNLP({"model_name": "bert"}, skip_model_load=True)
+
+    attention_mask = component._compute_attention_mask(
+        [actual_sequence_length], max_input_sequence_length
+    )
+    mask_ones = attention_mask[0][:zero_start_index]
+    mask_zeros = attention_mask[0][zero_start_index:]
+
+    assert np.all(mask_ones == 1)
+    assert np.all(mask_zeros == 0)


Code review comments:
diff_hunk: @@ -81,3 +82,45 @@ def test_input_padding(
     if padding_added:
         original_length = len(token_ids[0])
         assert np.all(np.array(padded_input[0][original_length:]) == 0)
+
+
+@pytest.mark.parametrize(
+    "sequence_length, model_name, should_overflow",
+    [(1000, "bert", True), (256, "bert", False)],
+)
+@pytest.mark.skip_on_windows
+def test_log_longer_sequence(
+    sequence_length: int, model_name: Text, should_overflow: bool, caplog
+):
+    transformers_config = {"model_name": model_name}
+
+    transformers_nlp = HFTransformersNLP(transformers_config)
+
+    text = " ".join(["hi"] * sequence_length)
+    message = Message(text)
+
+    caplog.set_level(logging.DEBUG)
+    transformers_nlp.process(message)
+    if should_overflow:
+        print(caplog.text)
path: tests/nlu/utils/test_hf_transformers.py
text: @user1:
I guess this can be removed.

@author:
Why would i miss that 🤯 
Thanks!