Repository: RasaHQ/rasa
Pull request: Tokenizers don't split words into sub-words
**Proposed changes**:
To avoid the problem of our entity extractors predicting entity labels for just a part of the words, we introduced a cleaning method after the prediction was done. However, we should avoid the incorrect prediction in the first place. 
To achieve this we will not tokenize words into sub-words anymore. We take the mean feature vectors of the sub-words as the feature vector of the word. 

fixes https://github.com/RasaHQ/rasa/issues/5755
closes https://github.com/RasaHQ/research/issues/83

**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:
EntityExtractor in v1.10.0 yield wrong entity value for language without spaces
Rasa version:
v1.10.0
Python version:
3.7
Operating system:
Windows, Linux

I'm not sure if this should be a bug or feature request, but here I go.

**Description of Problem**:
For language that doesn't have space that separate words, any `EntityExtractor` class will return the whole sentence as the `value`. When I see what changes have been made into that class, I see that there are few new functions that make up into `clean_up_entities` function. The problem appears from that function as I observe the process from the `SpacyEntityExtractor`. Based on the comment on function `_token_clusters`: "two tokens belong to the same word if there is no other character between them", I assume that all the process of the `clean_up_entities` function will merge whole input if there is no space between words.

**Overview of the Solution**:
I haven't really looked into the sub-words problem that requires this `clean_up_entities` function, so I can't really offer any solution except that I will still use v1.9.x. However, if this function will still be kept as it is for the future release, I either need to remake all of my training data to have space and preprocess the input to Rasa chatbot with custom components, or still stay on hold with version 1.9.x

**Examples**
On previous version:
```
>>> model.parse('木村拓哉は誰ですか')
'entities': [{'entity': 'Person',
   'value': '木村拓哉',
   'start': 0,
   'confidence': None,
   'end': 4,
   'extractor': 'SpacyEntityExtractor'}]
```
on version 1.10.0

```
>>> model.parse('木村拓哉は誰ですか')
'entities': [{'entity': 'Person',
   'value': '木村拓哉は誰ですか',
   'start': 0,
   'confidence': None,
   'end': 9,
   'extractor': 'SpacyEntityExtractor'}]
```

Hints from issue discussion:
Hello @tabergma,
Thank you very much for the quick response and I really appreciate it that you're already working on it.

Commit message: add test for HFTransfomerNLP

Patch to review:
diff --git a/changelog/5756.improvement.rst b/changelog/5756.improvement.rst
new file mode 100644
index 000000000000..d2a77898a7ad
--- /dev/null
+++ b/changelog/5756.improvement.rst
@@ -0,0 +1,5 @@
+To avoid the problem of our entity extractors predicting entity labels for just a part of the words,
+we introduced a cleaning method after the prediction was done.
+We should avoid the incorrect prediction in the first place.
+To achieve this we will not tokenize words into sub-words anymore.
+We take the mean feature vectors of the sub-words as the feature vector of the word.
\ No newline at end of file
diff --git a/rasa/nlu/classifiers/diet_classifier.py b/rasa/nlu/classifiers/diet_classifier.py
index 0f2ee5543958..32dfabe7cbb6 100644
--- a/rasa/nlu/classifiers/diet_classifier.py
+++ b/rasa/nlu/classifiers/diet_classifier.py
@@ -803,7 +803,6 @@ def _predict_entities(
         )
 
         entities = self.add_extractor_name(entities)
-        entities = self.clean_up_entities(message, entities)
         entities = message.get(ENTITIES, []) + entities
 
         return entities
diff --git a/rasa/nlu/constants.py b/rasa/nlu/constants.py
index 2e32e85e4acb..1ab586af0bf5 100644
--- a/rasa/nlu/constants.py
+++ b/rasa/nlu/constants.py
@@ -37,6 +37,7 @@
 
 CLS_TOKEN = "__CLS__"
 POSITION_OF_CLS_TOKEN = -1
+NUMBER_OF_SUB_TOKENS = "number_of_sub_tokens"
 
 MESSAGE_ATTRIBUTES = [TEXT, INTENT, RESPONSE]
 
diff --git a/rasa/nlu/extractors/crf_entity_extractor.py b/rasa/nlu/extractors/crf_entity_extractor.py
index d2a9c04c33fe..2d9b09967723 100644
--- a/rasa/nlu/extractors/crf_entity_extractor.py
+++ b/rasa/nlu/extractors/crf_entity_extractor.py
@@ -195,7 +195,6 @@ def _update_crf_order(self, training_data: TrainingData):
     def process(self, message: Message, **kwargs: Any) -> None:
         entities = self.extract_entities(message)
         entities = self.add_extractor_name(entities)
-        entities = self.clean_up_entities(message, entities)
         message.set(ENTITIES, message.get(ENTITIES, []) + entities, add_to_output=True)
 
     def extract_entities(self, message: Message) -> List[Dict[Text, Any]]:
diff --git a/rasa/nlu/extractors/duckling_http_extractor.py b/rasa/nlu/extractors/duckling_http_extractor.py
index 6033fd4dddbc..dff0f070c272 100644
--- a/rasa/nlu/extractors/duckling_http_extractor.py
+++ b/rasa/nlu/extractors/duckling_http_extractor.py
@@ -186,7 +186,6 @@ def process(self, message: Message, **kwargs: Any) -> None:
             )
 
         extracted = self.add_extractor_name(extracted)
-        extracted = self.clean_up_entities(message, extracted)
         message.set(ENTITIES, message.get(ENTITIES, []) + extracted, add_to_output=True)
 
     @classmethod
diff --git a/rasa/nlu/extractors/extractor.py b/rasa/nlu/extractors/extractor.py
index 983fd99e6669..7be4ef7082dc 100644
--- a/rasa/nlu/extractors/extractor.py
+++ b/rasa/nlu/extractors/extractor.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Text, Tuple, Optional, Union
+from typing import Any, Dict, List, Text, Tuple, Optional
 
 from rasa.constants import DOCS_URL_TRAINING_DATA_NLU
 from rasa.nlu.training_data import TrainingData
@@ -41,265 +41,6 @@ def add_processor_name(self, entity: Dict[Text, Any]) -> Dict[Text, Any]:
 
         return entity
 
-    def clean_up_entities(
-        self, message: Message, entities: List[Dict[Text, Any]], keep: bool = True
-    ) -> List[Dict[Text, Any]]:
-        """
-        Check if multiple entity labels are assigned to one word or if an entity label
-        is assigned to just a part of a word or if an entity label covers multiple
-        words, but one word just partly.
-
-        This might happen if you are using a tokenizer that splits up words into
-        sub-words and different entity labels are assigned to the individual sub-words.
-        If multiple entity labels are assigned to one word, we keep the entity label
-        with the highest confidence as entity label for that word. If just a part
-        of the word is annotated, that entity label is taken for the complete word.
-        If you set 'keep' to 'False', all entity labels for the word will be removed.
-
-        Args:
-            message: message object
-            entities: list of entities
-            keep:
-                If set to 'True', the entity label with the highest confidence is kept
-                if multiple entity labels are assigned to one word. If set to 'False'
-                all entity labels for that word will be removed.
-
-        Returns:
-            Updated entities.
-        """
-        misaligned_entities = self._get_misaligned_entities(
-            message.get(TOKENS_NAMES[TEXT]), entities
-        )
-
-        entity_indices_to_remove = set()
-
-        for misaligned_entity in misaligned_entities:
-            # entity indices involved in the misalignment
-            entity_indices = misaligned_entity["entity_indices"]
-
-            if not keep:
-                entity_indices_to_remove.update(entity_indices)
-                continue
-
-            idx = self._entity_index_to_keep(entities, entity_indices)
-
-            if idx is None or idx not in entity_indices:
-                entity_indices_to_remove.update(entity_indices)
-            else:
-                # keep just one entity
-                entity_indices.remove(idx)
-                entity_indices_to_remove.update(entity_indices)
-
-                # update that entity to cover the complete word(s)
-                entities[idx][ENTITY_ATTRIBUTE_START] = misaligned_entity[
-                    ENTITY_ATTRIBUTE_START
-                ]
-                entities[idx][ENTITY_ATTRIBUTE_END] = misaligned_entity[
-                    ENTITY_ATTRIBUTE_END
-                ]
-                entities[idx][ENTITY_ATTRIBUTE_VALUE] = message.text[
-                    misaligned_entity[ENTITY_ATTRIBUTE_START] : misaligned_entity[
-                        ENTITY_ATTRIBUTE_END
-                    ]
-                ]
-
-        # sort indices to remove entries at the end of the list first
-        # to avoid index out of range errors
-        for idx in sorted(entity_indices_to_remove, reverse=True):
-            entities.remove(entities[idx])
-
-        return entities
-
-    def _get_misaligned_entities(
-        self, tokens: List[Token], entities: List[Dict[Text, Any]]
-    ) -> List[Dict[Text, Any]]:
-        """Identify entities and tokens that are misaligned.
-
-        Misaligned entities are those that apply only to a part of a word, i.e.
-        sub-word.
-
-        Args:
-            tokens: list of tokens
-            entities: list of detected entities by the entity extractor
-
-        Returns:
-            Misaligned entities including the start and end position
-            of the final entity in the text and entity indices that are part of this
-            misalignment.
-        """
-        if not tokens:
-            return []
-
-        # group tokens: one token cluster corresponds to one word
-        token_clusters = self._token_clusters(tokens)
-
-        # added for tests, should only happen if tokens are not set or len(tokens) == 1
-        if not token_clusters:
-            return []
-
-        misaligned_entities = []
-        for entity_idx, entity in enumerate(entities):
-            # get all tokens that are covered/touched by the entity
-            entity_tokens = self._tokens_of_entity(entity, token_clusters)
-
-            if len(entity_tokens) == 1:
-                # entity covers exactly one word
-                continue
-
-            # get start and end position of complete word
-            # needed to update the final entity later
-            start_position = entity_tokens[0].start
-            end_position = entity_tokens[-1].end
-
-            # check if an entity was already found that covers the exact same word(s)
-            _idx = self._misaligned_entity_index(
-                misaligned_entities, start_position, end_position
-            )
-
-            if _idx is None:
-                misaligned_entities.append(
-                    {
-                        ENTITY_ATTRIBUTE_START: start_position,
-                        ENTITY_ATTRIBUTE_END: end_position,
-                        "entity_indices": [entity_idx],
-                    }
-                )
-            else:
-                # pytype: disable=attribute-error
-                misaligned_entities[_idx]["entity_indices"].append(entity_idx)
-                # pytype: enable=attribute-error
-
-        return misaligned_entities
-
-    @staticmethod
-    def _misaligned_entity_index(
-        word_entity_cluster: List[Dict[Text, Union[int, List[int]]]],
-        start_position: int,
-        end_position: int,
-    ) -> Optional[int]:
-        """Get index of matching misaligned entity.
-
-        Args:
-            word_entity_cluster: word entity cluster
-            start_position: start position
-            end_position: end position
-
-        Returns:
-            Index of the misaligned entity that matches the provided start and end
-            position.
-        """
-        for idx, cluster in enumerate(word_entity_cluster):
-            if (
-                cluster[ENTITY_ATTRIBUTE_START] == start_position
-                and cluster[ENTITY_ATTRIBUTE_END] == end_position
-            ):
-                return idx
-        return None
-
-    @staticmethod
-    def _tokens_of_entity(
-        entity: Dict[Text, Any], token_clusters: List[List[Token]]
-    ) -> List[Token]:
-        """Get all tokens of token clusters that are covered by the entity.
-
-        The entity can cover them completely or just partly.
-
-        Args:
-            entity: the entity
-            token_clusters: list of token clusters
-
-        Returns:
-            Token clusters that belong to the provided entity.
-
-        """
-        entity_tokens = []
-        for token_cluster in token_clusters:
-            entity_starts_inside_cluster = (
-                token_cluster[0].start
-                <= entity[ENTITY_ATTRIBUTE_START]
-                <= token_cluster[-1].end
-            )
-            entity_ends_inside_cluster = (
-                token_cluster[0].start
-                <= entity[ENTITY_ATTRIBUTE_END]
-                <= token_cluster[-1].end
-            )
-
-            if entity_starts_inside_cluster or entity_ends_inside_cluster:
-                entity_tokens += token_cluster
-        return entity_tokens
-
-    @staticmethod
-    def _token_clusters(tokens: List[Token]) -> List[List[Token]]:
-        """Build clusters of tokens that belong to one word.
-
-        Args:
-            tokens: list of tokens
-
-        Returns:
-            Token clusters.
-
-        """
-        # token cluster = list of token indices that belong to one word
-        token_index_clusters = []
-
-        # start at 1 in order to check if current token and previous token belong
-        # to the same word
-        for token_idx in range(1, len(tokens)):
-            previous_token_idx = token_idx - 1
-            # two tokens belong to the same word if there is no other character
-            # between them
-            if tokens[token_idx].start == tokens[previous_token_idx].end:
-                # a word was split into multiple tokens
-                token_cluster_already_exists = (
-                    token_index_clusters
-                    and token_index_clusters[-1][-1] == previous_token_idx
-                )
-                if token_cluster_already_exists:
-                    token_index_clusters[-1].append(token_idx)
-                else:
-                    token_index_clusters.append([previous_token_idx, token_idx])
-            else:
-                # the token corresponds to a single word
-                if token_idx == 1:
-                    token_index_clusters.append([previous_token_idx])
-                token_index_clusters.append([token_idx])
-
-        return [[tokens[idx] for idx in cluster] for cluster in token_index_clusters]
-
-    @staticmethod
-    def _entity_index_to_keep(
-        entities: List[Dict[Text, Any]], entity_indices: List[int]
-    ) -> Optional[int]:
-        """
-        Determine the entity index to keep.
-
-        If we just have one entity index, i.e. candidate, we return the index of that
-        candidate. If we have multiple candidates, we return the index of the entity
-        value with the highest confidence score. If no confidence score is present,
-        no entity label will be kept.
-
-        Args:
-            entities: the full list of entities
-            entity_indices: the entity indices to consider
-
-        Returns: the idx of the entity to keep
-        """
-        if len(entity_indices) == 1:
-            return entity_indices[0]
-
-        confidences = [
-            entities[idx][ENTITY_ATTRIBUTE_CONFIDENCE_TYPE]
-            for idx in entity_indices
-            if ENTITY_ATTRIBUTE_CONFIDENCE_TYPE in entities[idx]
-        ]
-
-        # we don't have confidence values for all entity labels
-        if len(confidences) != len(entity_indices):
-            return None
-
-        return confidences.index(max(confidences))
-
     @staticmethod
     def filter_irrelevant_entities(extracted: list, requested_dimensions: set) -> list:
         """Only return dimensions the user configured"""
diff --git a/rasa/nlu/extractors/mitie_entity_extractor.py b/rasa/nlu/extractors/mitie_entity_extractor.py
index 12ebbae1bbd1..140bb65ee2e8 100644
--- a/rasa/nlu/extractors/mitie_entity_extractor.py
+++ b/rasa/nlu/extractors/mitie_entity_extractor.py
@@ -141,7 +141,6 @@ def process(self, message: Message, **kwargs: Any) -> None:
             mitie_feature_extractor,
         )
         extracted = self.add_extractor_name(ents)
-        extracted = self.clean_up_entities(message, extracted)
         message.set(ENTITIES, message.get(ENTITIES, []) + extracted, add_to_output=True)
 
     @classmethod
diff --git a/rasa/nlu/extractors/spacy_entity_extractor.py b/rasa/nlu/extractors/spacy_entity_extractor.py
index f43395a95f75..d0546e031b8f 100644
--- a/rasa/nlu/extractors/spacy_entity_extractor.py
+++ b/rasa/nlu/extractors/spacy_entity_extractor.py
@@ -32,7 +32,6 @@ def process(self, message: Message, **kwargs: Any) -> None:
         spacy_nlp = kwargs.get("spacy_nlp", None)
         doc = spacy_nlp(message.text)
         all_extracted = self.add_extractor_name(self.extract_entities(doc))
-        all_extracted = self.clean_up_entities(message, all_extracted)
         dimensions = self.component_config["dimensions"]
         extracted = SpacyEntityExtractor.filter_irrelevant_entities(
             all_extracted, dimensions
diff --git a/rasa/nlu/featurizers/dense_featurizer/convert_featurizer.py b/rasa/nlu/featurizers/dense_featurizer/convert_featurizer.py
index 1fa9de8d3210..9848149d679b 100644
--- a/rasa/nlu/featurizers/dense_featurizer/convert_featurizer.py
+++ b/rasa/nlu/featurizers/dense_featurizer/convert_featurizer.py
@@ -9,12 +9,7 @@
 from rasa.nlu.tokenizers.convert_tokenizer import ConveRTTokenizer
 from rasa.nlu.config import RasaNLUModelConfig
 from rasa.nlu.training_data import Message, TrainingData
-from rasa.nlu.constants import (
-    TEXT,
-    TOKENS_NAMES,
-    DENSE_FEATURE_NAMES,
-    DENSE_FEATURIZABLE_ATTRIBUTES,
-)
+from rasa.nlu.constants import TEXT, DENSE_FEATURE_NAMES, DENSE_FEATURIZABLE_ATTRIBUTES
 import numpy as np
 import tensorflow as tf
 
@@ -79,25 +74,28 @@ def _compute_sequence_encodings(
         self, batch_examples: List[Message], attribute: Text = TEXT
     ) -> Tuple[np.ndarray, List[int]]:
         list_of_tokens = [
-            example.get(TOKENS_NAMES[attribute]) for example in batch_examples
+            train_utils.tokens_without_cls(example, attribute)
+            for example in batch_examples
         ]
 
-        # remove CLS token from list of tokens
-        list_of_tokens = [sent_tokens[:-1] for sent_tokens in list_of_tokens]
-
         number_of_tokens_in_sentence = [
             len(sent_tokens) for sent_tokens in list_of_tokens
         ]
 
         # join the tokens to get a clean text to ensure the sequence length of
         # the returned embeddings from ConveRT matches the length of the tokens
+        # (including sub-tokens)
         tokenized_texts = self._tokens_to_text(list_of_tokens)
+        token_features = self._sequence_encoding_of_text(tokenized_texts)
 
-        return (
-            self._sequence_encoding_of_text(tokenized_texts),
-            number_of_tokens_in_sentence,
+        # ConveRT might split up tokens into sub-tokens
+        # take the mean of the sub-token vectors and use that as the token vector
+        token_features = train_utils.align_token_features(
+            list_of_tokens, token_features
         )
 
+        return token_features, number_of_tokens_in_sentence
+
     def _combine_encodings(
         self,
         sentence_encodings: np.ndarray,
diff --git a/rasa/nlu/featurizers/featurizer.py b/rasa/nlu/featurizers/featurizer.py
index ef896a2a5ab8..0e36b4f85b50 100644
--- a/rasa/nlu/featurizers/featurizer.py
+++ b/rasa/nlu/featurizers/featurizer.py
@@ -1,6 +1,7 @@
 import numpy as np
 import scipy.sparse
 from typing import Any, Text, Union, Optional
+
 from rasa.nlu.training_data import Message
 from rasa.nlu.components import Component
 from rasa.nlu.constants import SPARSE_FEATURE_NAMES, DENSE_FEATURE_NAMES, TEXT
diff --git a/rasa/nlu/tokenizers/convert_tokenizer.py b/rasa/nlu/tokenizers/convert_tokenizer.py
index 3186e1150207..c73d681541ea 100644
--- a/rasa/nlu/tokenizers/convert_tokenizer.py
+++ b/rasa/nlu/tokenizers/convert_tokenizer.py
@@ -1,5 +1,6 @@
 from typing import Any, Dict, List, Text
 
+from rasa.nlu.constants import NUMBER_OF_SUB_TOKENS
 from rasa.nlu.tokenizers.tokenizer import Token
 from rasa.nlu.tokenizers.whitespace_tokenizer import WhitespaceTokenizer
 from rasa.nlu.training_data import Message
@@ -46,8 +47,8 @@ def tokenize(self, message: Message, attribute: Text) -> List[Token]:
         ConveRT adds a special char in front of (some) words and splits words into
         sub-words. To ensure the entity start and end values matches the token values,
         tokenize the text first using the whitespace tokenizer. If individual tokens
-        are split up into multiple tokens, we make sure that the start end end value
-        of the first and last respective tokens stay the same.
+        are split up into multiple tokens, add this information to the
+        respected tokens.
         """
 
         # perform whitespace tokenization
@@ -56,17 +57,15 @@ def tokenize(self, message: Message, attribute: Text) -> List[Token]:
         tokens_out = []
 
         for token in tokens_in:
-            token_start, token_end, token_text = token.start, token.end, token.text
-
             # use ConveRT model to tokenize the text
-            split_token_strings = self._tokenize(token_text)[0]
+            split_token_strings = self._tokenize(token.text)[0]
 
             # clean tokens (remove special chars and empty tokens)
             split_token_strings = self._clean_tokens(split_token_strings)
 
-            tokens_out += train_utils.align_tokens(
-                split_token_strings, token_end, token_start
-            )
+            token.set(NUMBER_OF_SUB_TOKENS, len(split_token_strings))
+
+            tokens_out.append(token)
 
         return tokens_out
 
diff --git a/rasa/nlu/utils/hugging_face/hf_transformers.py b/rasa/nlu/utils/hugging_face/hf_transformers.py
index a40d926bce0c..889bf481356d 100644
--- a/rasa/nlu/utils/hugging_face/hf_transformers.py
+++ b/rasa/nlu/utils/hugging_face/hf_transformers.py
@@ -17,6 +17,7 @@
     TOKENS,
     SENTENCE_FEATURES,
     SEQUENCE_FEATURES,
+    NUMBER_OF_SUB_TOKENS,
 )
 
 logger = logging.getLogger(__name__)
@@ -116,7 +117,8 @@ 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.
@@ -139,12 +141,15 @@ def _lm_specific_token_cleanup(
     ) -> Tuple[List[int], List[Text]]:
         """Clean up special chars added by tokenizers of language models.
 
-        Many language models add a special char in front/back of (some) words. We clean up those chars as they are not
+        Many language models add a special char in front/back of (some) words. We clean
+        up those chars as they are not
         needed once the features are already computed.
 
         Args:
-            split_token_ids: List of token ids received as output from the language model specific tokenizer.
-            token_strings: List of token strings received as output from the language model specific tokenizer.
+            split_token_ids: List of token ids received as output from the language
+            model specific tokenizer.
+            token_strings: List of token strings received as output from the language
+            model specific tokenizer.
 
         Returns:
             Cleaned up token ids and token strings.
@@ -156,10 +161,12 @@ def _lm_specific_token_cleanup(
     def _post_process_sequence_embeddings(
         self, sequence_embeddings: np.ndarray
     ) -> Tuple[np.ndarray, np.ndarray]:
-        """Compute sentence level representations and sequence level representations for relevant tokens.
+        """Compute sentence level representations and sequence level representations
+        for relevant tokens.
 
         Args:
-            sequence_embeddings: Sequence level dense features received as output from language model.
+            sequence_embeddings: Sequence level dense features received as output from
+            language model.
 
         Returns:
             Sentence and sequence level representations.
@@ -191,18 +198,20 @@ def _tokenize_example(
     ) -> Tuple[List[Token], List[int]]:
         """Tokenize a single message example.
 
-        Many language models add a special char in front of (some) words and split words into
-        sub-words. To ensure the entity start and end values matches the token values,
-        tokenize the text first using the whitespace tokenizer. If individual tokens
-        are split up into multiple tokens, we make sure that the start and end value
-        of the first and last respective tokens stay the same.
+        Many language models add a special char in front of (some) words and split
+        words into sub-words. To ensure the entity start and end values matches the
+        token values, tokenize the text first using the whitespace tokenizer. If
+        individual tokens are split up into multiple tokens, we add this information
+        to the respected token.
 
         Args:
             message: Single message object to be processed.
-            attribute: Property of message to be processed, one of ``TEXT`` or ``RESPONSE``.
+            attribute: Property of message to be processed, one of ``TEXT`` or
+            ``RESPONSE``.
 
         Returns:
-            List of token strings and token ids for the corresponding attribute of the message.
+            List of token strings and token ids for the corresponding attribute of the
+            message.
         """
 
         tokens_in = self.whitespace_tokenizer.tokenize(message, attribute)
@@ -221,9 +230,9 @@ def _tokenize_example(
 
             token_ids_out += split_token_ids
 
-            tokens_out += train_utils.align_tokens(
-                split_token_strings, token.end, token.start
-            )
+            token.set(NUMBER_OF_SUB_TOKENS, len(split_token_strings))
+
+            tokens_out.append(token)
 
         return tokens_out, token_ids_out
 
@@ -234,8 +243,10 @@ def _get_token_ids_for_batch(
 
         A token id is the id of that token in the vocabulary of the language model.
         Args:
-            batch_examples: Batch of message objects for which tokens need to be computed.
-            attribute: Property of message to be processed, one of ``TEXT`` or ``RESPONSE``.
+            batch_examples: Batch of message objects for which tokens need to be
+            computed.
+            attribute: Property of message to be processed, one of ``TEXT`` or
+            ``RESPONSE``.
 
         Returns:
             List of token strings and token ids for each example in the batch.
@@ -257,7 +268,8 @@ def _get_token_ids_for_batch(
     def _compute_attention_mask(actual_sequence_lengths: List[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.
+        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
@@ -286,7 +298,8 @@ def _add_padding_to_batch(
         """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.
+            batch_token_ids: Batch of examples where each example is a non-padded list
+            of token ids.
 
         Returns:
             Padded batch with all examples of the same length.
@@ -315,7 +328,8 @@ def _add_padding_to_batch(
     def _extract_nonpadded_embeddings(
         embeddings: np.ndarray, actual_sequence_lengths: List[int]
     ) -> np.ndarray:
-        """Use pre-computed non-padded lengths of each example to extract embeddings for non-padding tokens.
+        """Use pre-computed non-padded lengths of each example to extract embeddings
+        for non-padding tokens.
 
         Args:
             embeddings: sequence level representations for each example of the batch.
@@ -337,8 +351,10 @@ def _compute_batch_sequence_features(
         """Feed the padded batch to the language model.
 
         Args:
-            batch_attention_mask: Mask of 0s and 1s which indicate whether the token is a padding token or not.
-            padded_token_ids: Batch of token ids for each example. The batch is padded and hence can be fed at once.
+            batch_attention_mask: Mask of 0s and 1s which indicate whether the token
+            is a padding token or not.
+            padded_token_ids: Batch of token ids for each example. The batch is padded
+            and hence can be fed at once.
 
         Returns:
             Sequence level representations from the language model.
@@ -354,14 +370,16 @@ def _compute_batch_sequence_features(
         return sequence_hidden_states
 
     def _get_model_features_for_batch(
-        self, batch_token_ids: List[List[int]]
+        self, batch_token_ids: List[List[int]], batch_tokens: List[List[Token]]
     ) -> Tuple[np.ndarray, np.ndarray]:
         """Compute dense features of each example in the batch.
 
-        We first add the special tokens corresponding to each language model. Next, we add appropriate padding
-        and compute a mask for that padding so that it doesn't affect the feature computation. The padded batch is next
-        fed to the language model and token level embeddings are computed. Using the pre-computed mask, embeddings for
-        non-padding tokens are extracted and subsequently sentence level embeddings are computed.
+        We first add the special tokens corresponding to each language model. Next, we
+        add appropriate padding and compute a mask for that padding so that it doesn't
+        affect the feature computation. The padded batch is next fed to the language
+        model and token level embeddings are computed. Using the pre-computed mask,
+        embeddings for non-padding tokens are extracted and subsequently sentence
+        level embeddings are computed.
 
         Args:
             batch_token_ids: List of token ids of each example in the batch.
@@ -395,9 +413,28 @@ def _get_model_features_for_batch(
         # Extract sentence level and post-processed features
         (
             sentence_embeddings,
-            sequence_final_embeddings,
+            sequence_embeddings,
         ) = self._post_process_sequence_embeddings(sequence_nonpadded_embeddings)
 
+        # shape of matrix for all sequence embeddings
+        batch_dim = len(sequence_embeddings)
+        seq_dim = max(e.shape[0] for e in sequence_embeddings)
+        feature_dim = sequence_embeddings[0].shape[1]
+        shape = (batch_dim, seq_dim, feature_dim)
+
+        # align features with tokens so that we have just one vector per token
+        # (don't include sub-tokens)
+        sequence_embeddings = train_utils.align_token_features(
+            batch_tokens, sequence_embeddings, shape
+        )
+
+        # sequence_embeddings is a padded numpy array
+        # remove the padding, keep just the non-zero vectors
+        sequence_final_embeddings = []
+        for embeddings, tokens in zip(sequence_embeddings, batch_tokens):
+            sequence_final_embeddings.append(embeddings[: len(tokens)])
+        sequence_final_embeddings = np.array(sequence_final_embeddings)
+
         return sentence_embeddings, sequence_final_embeddings
 
     def _get_docs_for_batch(
@@ -406,8 +443,10 @@ def _get_docs_for_batch(
         """Compute language model docs for all examples in the batch.
 
         Args:
-            batch_examples: Batch of message objects for which language model docs need to be computed.
-            attribute: Property of message to be processed, one of ``TEXT`` or ``RESPONSE``.
+            batch_examples: Batch of message objects for which language model docs
+            need to be computed.
+            attribute: Property of message to be processed, one of ``TEXT`` or
+            ``RESPONSE``.
 
         Returns:
             List of language model docs for each message in batch.
@@ -420,10 +459,11 @@ def _get_docs_for_batch(
         (
             batch_sentence_features,
             batch_sequence_features,
-        ) = self._get_model_features_for_batch(batch_token_ids)
+        ) = self._get_model_features_for_batch(batch_token_ids, batch_tokens)
 
         # A doc consists of
-        # {'token_ids': ..., 'tokens': ..., 'sequence_features': ..., 'sentence_features': ...}
+        # {'token_ids': ..., 'tokens': ..., 'sequence_features': ...,
+        # 'sentence_features': ...}
         batch_docs = []
         for index in range(len(batch_examples)):
             doc = {
@@ -468,7 +508,8 @@ def train(
                 # Collect batch examples
                 batch_messages = non_empty_examples[batch_start_index:batch_end_index]
 
-                # Construct a doc with relevant features extracted(tokens, dense_features)
+                # Construct a doc with relevant features
+                # extracted(tokens, dense_features)
                 batch_docs = self._get_docs_for_batch(batch_messages, attribute)
 
                 for index, ex in enumerate(batch_messages):
diff --git a/rasa/utils/train_utils.py b/rasa/utils/train_utils.py
index efec99c08487..aaa3fd2767c6 100644
--- a/rasa/utils/train_utils.py
+++ b/rasa/utils/train_utils.py
@@ -1,7 +1,7 @@
 import numpy as np
 import logging
 import scipy.sparse
-from typing import Optional, Text, Dict, Any, Union, List
+from typing import Optional, Text, Dict, Any, Union, List, Tuple
 
 from rasa.nlu.training_data import Message
 from rasa.core.constants import DIALOGUE
@@ -10,6 +10,7 @@
     TOKENS_NAMES,
     DENSE_FEATURIZABLE_ATTRIBUTES,
     POSITION_OF_CLS_TOKEN,
+    NUMBER_OF_SUB_TOKENS,
 )
 from rasa.nlu.tokenizers.tokenizer import Token
 import rasa.utils.io as io_utils
@@ -79,40 +80,51 @@ def update_similarity_type(config: Dict[Text, Any]) -> Dict[Text, Any]:
     return config
 
 
-def align_tokens(
-    tokens_in: List[Text], token_end: int, token_start: int
-) -> List[Token]:
-    """Align sub-tokens of Language model with tokens return by the WhitespaceTokenizer.
+def align_token_features(
+    list_of_tokens: List[List[Token]],
+    in_token_features: np.ndarray,
+    shape: Optional[Tuple] = None,
+) -> np.ndarray:
+    """Align token features to match tokens.
 
-    As a language model might split a single word into multiple tokens, we need to make
-    sure that the start and end value of first and last sub-token matches the
-    start and end value of the token return by the WhitespaceTokenizer as the
-    entities are using those start and end values.
-    """
-
-    tokens_out = []
+    ConveRT might split up tokens into sub-tokens. We need to take the mean of
+    the sub-token vectors and take that as token vector.
 
-    current_token_offset = token_start
+    Args:
+        list_of_tokens: tokens for examples
+        in_token_features: token features from ConveRT
+        shape: shape of feature matrix
 
-    for index, string in enumerate(tokens_in):
-        if index == 0:
-            if index == len(tokens_in) - 1:
-                s_token_end = token_end
-            else:
-                s_token_end = current_token_offset + len(string)
-            tokens_out.append(Token(string, token_start, end=s_token_end))
-        elif index == len(tokens_in) - 1:
-            tokens_out.append(Token(string, current_token_offset, end=token_end))
-        else:
-            tokens_out.append(
-                Token(
-                    string, current_token_offset, end=current_token_offset + len(string)
+    Returns:
+        Token features.
+    """
+    if shape is None:
+        shape = in_token_features.shape
+    out_token_features = np.zeros(shape)
+
+    for example_idx, example_tokens in enumerate(list_of_tokens):
+        offset = 0
+        for token_idx, token in enumerate(example_tokens):
+            number_sub_words = token.get(NUMBER_OF_SUB_TOKENS, 1)
+
+            if number_sub_words > 1:
+                token_start_idx = token_idx + offset
+                token_end_idx = token_idx + offset + number_sub_words
+
+                mean_vec = np.mean(
+                    in_token_features[example_idx][token_start_idx:token_end_idx],
+                    axis=0,
                 )
-            )
 
-        current_token_offset += len(string)
+                offset += number_sub_words - 1
+
+                out_token_features[example_idx][token_idx] = mean_vec
+            else:
+                out_token_features[example_idx][token_idx] = in_token_features[
+                    example_idx
+                ][token_idx + offset]
 
-    return tokens_out
+    return out_token_features
 
 
 def sequence_to_sentence_features(
diff --git a/tests/nlu/extractors/test_extractor.py b/tests/nlu/extractors/test_extractor.py
index 119f569b4b56..d5a2f5852d00 100644
--- a/tests/nlu/extractors/test_extractor.py
+++ b/tests/nlu/extractors/test_extractor.py
@@ -3,253 +3,12 @@
 import pytest
 
 from rasa.nlu.constants import TEXT
-from rasa.nlu.tokenizers.tokenizer import Token
 from rasa.nlu.training_data import Message
 from rasa.nlu.extractors.extractor import EntityExtractor
 from rasa.nlu.tokenizers.whitespace_tokenizer import WhitespaceTokenizer
 from rasa.nlu.training_data.formats import MarkdownReader
 
 
-@pytest.mark.parametrize(
-    "text, tokens, entities, keep, expected_entities",
-    [
-        (
-            "Aarhus is a city",
-            [
-                Token("Aar", 0, 3),
-                Token("hus", 3, 6),
-                Token("is", 7, 9),
-                Token("a", 10, 11),
-                Token("city", 12, 16),
-            ],
-            [
-                {"entity": "iata", "start": 0, "end": 3, "value": "Aar"},
-                {"entity": "city", "start": 3, "end": 6, "value": "hus"},
-                {"entity": "location", "start": 12, "end": 16, "value": "city"},
-            ],
-            False,
-            [{"entity": "location", "start": 12, "end": 16, "value": "city"}],
-        ),
-        (
-            "Aarhus",
-            [Token("Aar", 0, 3), Token("hus", 3, 6)],
-            [
-                {"entity": "iata", "start": 0, "end": 3, "value": "Aar"},
-                {"entity": "city", "start": 3, "end": 6, "value": "hus"},
-            ],
-            True,
-            [],
-        ),
-        (
-            "Aarhus city",
-            [Token("Aarhus", 0, 6), Token("city", 7, 11)],
-            [
-                {"entity": "city", "start": 0, "end": 6, "value": "Aarhus"},
-                {"entity": "type", "start": 7, "end": 11, "value": "city"},
-            ],
-            False,
-            [
-                {"entity": "city", "start": 0, "end": 6, "value": "Aarhus"},
-                {"entity": "type", "start": 7, "end": 11, "value": "city"},
-            ],
-        ),
-        (
-            "Aarhus is a city",
-            [
-                Token("Aar", 0, 3),
-                Token("hus", 3, 6),
-                Token("is", 7, 9),
-                Token("a", 10, 11),
-                Token("city", 12, 16),
-            ],
-            [
-                {
-                    "entity": "city",
-                    "start": 0,
-                    "end": 3,
-                    "confidence_entity": 0.87,
-                    "value": "Aar",
-                },
-                {
-                    "entity": "iata",
-                    "start": 3,
-                    "end": 6,
-                    "confidence_entity": 0.43,
-                    "value": "hus",
-                },
-                {"entity": "location", "start": 12, "end": 16, "value": "city"},
-            ],
-            True,
-            [
-                {
-                    "entity": "city",
-                    "start": 0,
-                    "end": 6,
-                    "confidence_entity": 0.87,
-                    "value": "Aarhus",
-                },
-                {"entity": "location", "start": 12, "end": 16, "value": "city"},
-            ],
-        ),
-        (
-            "Aarhus",
-            [Token("Aa", 0, 2), Token("r", 2, 3), Token("hu", 3, 5), Token("s", 5, 6)],
-            [
-                {
-                    "entity": "iata",
-                    "start": 0,
-                    "end": 2,
-                    "confidence_entity": 0.32,
-                    "value": "Aa",
-                },
-                {
-                    "entity": "city",
-                    "start": 2,
-                    "end": 3,
-                    "confidence_entity": 0.87,
-                    "value": "r",
-                },
-                {
-                    "entity": "iata",
-                    "start": 3,
-                    "end": 5,
-                    "confidence_entity": 0.21,
-                    "value": "hu",
-                },
-                {
-                    "entity": "city",
-                    "start": 5,
-                    "end": 6,
-                    "confidence_entity": 0.43,
-                    "value": "s",
-                },
-            ],
-            True,
-            [
-                {
-                    "entity": "city",
-                    "start": 0,
-                    "end": 6,
-                    "confidence_entity": 0.87,
-                    "value": "Aarhus",
-                }
-            ],
-        ),
-        (
-            "Aarhus",
-            [Token("Aa", 0, 2), Token("r", 2, 3), Token("hu", 3, 5), Token("s", 5, 6)],
-            [
-                {
-                    "entity": "city",
-                    "start": 0,
-                    "end": 2,
-                    "confidence_entity": 0.32,
-                    "value": "Aa",
-                }
-            ],
-            True,
-            [
-                {
-                    "entity": "city",
-                    "start": 0,
-                    "end": 6,
-                    "confidence_entity": 0.32,
-                    "value": "Aarhus",
-                }
-            ],
-        ),
-        (
-            "Aarhus",
-            [Token("Aa", 0, 2), Token("r", 2, 3), Token("hu", 3, 5), Token("s", 5, 6)],
-            [
-                {
-                    "entity": "city",
-                    "start": 0,
-                    "end": 2,
-                    "confidence_entity": 0.32,
-                    "value": "Aa",
-                }
-            ],
-            False,
-            [],
-        ),
-        (
-            "Buenos Aires is a city",
-            [
-                Token("Buenos", 0, 6),
-                Token("Ai", 7, 9),
-                Token("res", 9, 12),
-                Token("is", 13, 15),
-                Token("a", 16, 17),
-                Token("city", 18, 22),
-            ],
-            [
-                {"entity": "city", "start": 0, "end": 9, "value": "Buenos Ai"},
-                {"entity": "location", "start": 18, "end": 22, "value": "city"},
-            ],
-            False,
-            [{"entity": "location", "start": 18, "end": 22, "value": "city"}],
-        ),
-        (
-            "Buenos Aires is a city",
-            [
-                Token("Buenos", 0, 6),
-                Token("Ai", 7, 9),
-                Token("res", 9, 12),
-                Token("is", 13, 15),
-                Token("a", 16, 17),
-                Token("city", 18, 22),
-            ],
-            [
-                {"entity": "city", "start": 0, "end": 9, "value": "Buenos Ai"},
-                {"entity": "location", "start": 18, "end": 22, "value": "city"},
-            ],
-            True,
-            [
-                {"entity": "city", "start": 0, "end": 12, "value": "Buenos Aires"},
-                {"entity": "location", "start": 18, "end": 22, "value": "city"},
-            ],
-        ),
-        (
-            "Buenos Aires is a city",
-            [
-                Token("Buen", 0, 4),
-                Token("os", 4, 6),
-                Token("Ai", 7, 9),
-                Token("res", 9, 12),
-                Token("is", 13, 15),
-                Token("a", 16, 17),
-                Token("city", 18, 22),
-            ],
-            [
-                {"entity": "city", "start": 4, "end": 9, "value": "os Ai"},
-                {"entity": "location", "start": 18, "end": 22, "value": "city"},
-            ],
-            True,
-            [
-                {"entity": "city", "start": 0, "end": 12, "value": "Buenos Aires"},
-                {"entity": "location", "start": 18, "end": 22, "value": "city"},
-            ],
-        ),
-    ],
-)
-def test_clean_up_entities(
-    text: Text,
-    tokens: List[Token],
-    entities: List[Dict[Text, Any]],
-    keep: bool,
-    expected_entities: List[Dict[Text, Any]],
-):
-    extractor = EntityExtractor()
-
-    message = Message(text)
-    message.set("tokens", tokens)
-
-    updated_entities = extractor.clean_up_entities(message, entities, keep)
-
-    assert updated_entities == expected_entities
-
-
 @pytest.mark.parametrize(
     "text, tags, confidences, expected_entities",
     [
diff --git a/tests/nlu/featurizers/test_convert_featurizer.py b/tests/nlu/featurizers/test_convert_featurizer.py
index c9dc0f47b3ba..e42ac68fde15 100644
--- a/tests/nlu/featurizers/test_convert_featurizer.py
+++ b/tests/nlu/featurizers/test_convert_featurizer.py
@@ -4,13 +4,7 @@
 from rasa.nlu.tokenizers.tokenizer import Tokenizer
 from rasa.nlu.training_data import TrainingData
 from rasa.nlu.tokenizers.convert_tokenizer import ConveRTTokenizer
-from rasa.nlu.constants import (
-    TEXT,
-    DENSE_FEATURE_NAMES,
-    TOKENS_NAMES,
-    RESPONSE,
-    INTENT,
-)
+from rasa.nlu.constants import TEXT, DENSE_FEATURE_NAMES, TOKENS_NAMES, RESPONSE, INTENT
 from rasa.nlu.training_data import Message
 from rasa.nlu.config import RasaNLUModelConfig
 from rasa.nlu.featurizers.dense_featurizer.convert_featurizer import ConveRTFeaturizer
diff --git a/tests/nlu/featurizers/test_lm_featurizer.py b/tests/nlu/featurizers/test_lm_featurizer.py
index f490781ad714..ef04edd7cef8 100644
--- a/tests/nlu/featurizers/test_lm_featurizer.py
+++ b/tests/nlu/featurizers/test_lm_featurizer.py
@@ -14,7 +14,7 @@
         (
             "bert",
             ["Good evening.", "here is the sentence I want embeddings for."],
-            [(3, 768), (12, 768)],
+            [(3, 768), (9, 768)],
             [
                 [0.5727445, -0.16078179],
                 [-0.5485125, 0.09632876, -0.4278888, 0.11438395, 0.18316492],
@@ -27,7 +27,7 @@
         (
             "gpt",
             ["Good evening.", "here is the sentence I want embeddings for."],
-            [(3, 768), (10, 768)],
+            [(3, 768), (9, 768)],
             [
                 [-0.0630323737859726, 0.4029877185821533],
                 [
@@ -58,16 +58,10 @@
         (
             "gpt2",
             ["Good evening.", "here is the sentence I want embeddings for."],
-            [(4, 768), (12, 768)],
+            [(3, 768), (9, 768)],
             [
-                [-0.033827826380729675, -0.10971662402153015, 0.002244209870696068],
-                [
-                    -0.18434514105319977,
-                    -0.5386468768119812,
-                    -0.11122681945562363,
-                    -1.368929147720337,
-                    -0.5397579669952393,
-                ],
+                [-0.03382749, -0.05373593],
+                [-0.18434484, -0.5386464, -0.11122551, -0.95434338, 0.28311089],
             ],
             [
                 [
@@ -89,7 +83,7 @@
         (
             "xlnet",
             ["Good evening.", "here is the sentence I want embeddings for."],
-            [(3, 768), (11, 768)],
+            [(3, 768), (9, 768)],
             [
                 [1.7612367868423462, 2.5819129943847656],
                 [
@@ -120,7 +114,7 @@
         (
             "distilbert",
             ["Good evening.", "here is the sentence I want embeddings for."],
-            [(3, 768), (12, 768)],
+            [(3, 768), (9, 768)],
             [
                 [0.22866562008857727, -0.0575055330991745],
                 [
@@ -151,16 +145,10 @@
         (
             "roberta",
             ["Good evening.", "here is the sentence I want embeddings for."],
-            [(4, 768), (12, 768)],
+            [(3, 768), (9, 768)],
             [
-                [-0.309267520904541, 0.12365783751010895, 0.06769893318414688],
-                [
-                    0.02152823843061924,
-                    -0.08026768267154694,
-                    -0.10808645188808441,
-                    0.20090824365615845,
-                    0.04756045714020729,
-                ],
+                [-0.3092685, 0.09567838],
+                [0.02152853, -0.08026707, -0.1080862, 0.12423468, -0.05378958],
             ],
             [
                 [
diff --git a/tests/nlu/tokenizers/test_convert_tokenizer.py b/tests/nlu/tokenizers/test_convert_tokenizer.py
index 06cd3a9e7bad..e480784286e0 100644
--- a/tests/nlu/tokenizers/test_convert_tokenizer.py
+++ b/tests/nlu/tokenizers/test_convert_tokenizer.py
@@ -9,7 +9,7 @@
     "text, expected_tokens, expected_indices",
     [
         (
-            "Forecast for lunch",
+            "forecast for lunch",
             ["forecast", "for", "lunch"],
             [(0, 8), (9, 12), (13, 18)],
         ),
@@ -17,11 +17,7 @@
         ("you're", ["you", "re"], [(0, 3), (4, 6)]),
         ("r. n. b.", ["r", "n", "b"], [(0, 1), (3, 4), (6, 7)]),
         ("rock & roll", ["rock", "&", "roll"], [(0, 4), (5, 6), (7, 11)]),
-        (
-            "ńöñàśçií",
-            ["ń", "ö", "ñ", "à", "ś", "ç", "i", "í"],
-            [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7), (7, 8)],
-        ),
+        ("ńöñàśçií", ["ńöñàśçií"], [(0, 8)]),
     ],
 )
 def test_convert_tokenizer_edge_cases(text, expected_tokens, expected_indices):
diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py
new file mode 100644
index 000000000000..ea773d2b6698
--- /dev/null
+++ b/tests/utils/test_train_utils.py
@@ -0,0 +1,43 @@
+import numpy as np
+
+import rasa.utils.train_utils as train_utils
+from nlu.tokenizers.lm_tokenizer import LanguageModelTokenizer
+from rasa.nlu.utils.hugging_face.hf_transformers import HFTransformersNLP
+from rasa.nlu.constants import TEXT, NUMBER_OF_SUB_TOKENS
+from rasa.nlu.tokenizers.convert_tokenizer import ConveRTTokenizer
+from rasa.nlu.training_data import Message
+
+
+def test_align_token_features_convert():
+    tokens = ConveRTTokenizer().tokenize(Message("In Aarhus and Ahaus"), attribute=TEXT)
+    seq_dim = sum(t.get(NUMBER_OF_SUB_TOKENS) for t in tokens)
+    token_features = np.random.rand(1, seq_dim, 64)
+
+    actual_features = train_utils.align_token_features([tokens], token_features)
+
+    assert np.all(actual_features[0][0] == token_features[0][0])
+    # Aarhus is split into 'aar' and 'hus'
+    assert np.all(actual_features[0][1] == np.mean(token_features[0][1:3], axis=0))
+    assert np.all(actual_features[0][2] == token_features[0][3])
+    # Ahaus is split into 'aha' and 'us'
+    assert np.all(actual_features[0][3] == np.mean(token_features[0][4:6], axis=0))
+
+
+def test_align_token_features_bert():
+    message = Message("a sentence with embeddings")
+
+    hf_nlp = HFTransformersNLP()
+    hf_nlp.process(message)
+
+    tokens = LanguageModelTokenizer().tokenize(message, attribute=TEXT)
+
+    seq_dim = sum(t.get(NUMBER_OF_SUB_TOKENS) for t in tokens)
+    token_features = np.random.rand(1, seq_dim, 64)
+
+    actual_features = train_utils.align_token_features([tokens], token_features)
+
+    assert np.all(actual_features[0][0] == token_features[0][0])
+    assert np.all(actual_features[0][1] == token_features[0][1])
+    assert np.all(actual_features[0][2] == token_features[0][2])
+    # embeddings is split into 4 sub-tokens
+    assert np.all(actual_features[0][3] == np.mean(token_features[0][3:7], axis=0))


Code review comments:
diff_hunk: @@ -0,0 +1,43 @@
+import numpy as np
+
+import rasa.utils.train_utils as train_utils
+from nlu.tokenizers.lm_tokenizer import LanguageModelTokenizer
+from rasa.nlu.utils.hugging_face.hf_transformers import HFTransformersNLP
+from rasa.nlu.constants import TEXT, NUMBER_OF_SUB_TOKENS
+from rasa.nlu.tokenizers.convert_tokenizer import ConveRTTokenizer
+from rasa.nlu.training_data import Message
+
+
+def test_align_token_features_convert():
+    tokens = ConveRTTokenizer().tokenize(Message("In Aarhus and Ahaus"), attribute=TEXT)
+    seq_dim = sum(t.get(NUMBER_OF_SUB_TOKENS) for t in tokens)
+    token_features = np.random.rand(1, seq_dim, 64)
+
+    actual_features = train_utils.align_token_features([tokens], token_features)
+
+    assert np.all(actual_features[0][0] == token_features[0][0])
+    # Aarhus is split into 'aar' and 'hus'
+    assert np.all(actual_features[0][1] == np.mean(token_features[0][1:3], axis=0))
+    assert np.all(actual_features[0][2] == token_features[0][3])
+    # Ahaus is split into 'aha' and 'us'
+    assert np.all(actual_features[0][3] == np.mean(token_features[0][4:6], axis=0))
+
+
+def test_align_token_features_bert():
+    message = Message("a sentence with embeddings")
+
+    hf_nlp = HFTransformersNLP()
+    hf_nlp.process(message)
+
+    tokens = LanguageModelTokenizer().tokenize(message, attribute=TEXT)
+
+    seq_dim = sum(t.get(NUMBER_OF_SUB_TOKENS) for t in tokens)
+    token_features = np.random.rand(1, seq_dim, 64)
+
+    actual_features = train_utils.align_token_features([tokens], token_features)
+
+    assert np.all(actual_features[0][0] == token_features[0][0])
+    assert np.all(actual_features[0][1] == token_features[0][1])
+    assert np.all(actual_features[0][2] == token_features[0][2])
+    # embeddings is split into 4 sub-tokens
+    assert np.all(actual_features[0][3] == np.mean(token_features[0][3:7], axis=0))
path: tests/utils/test_train_utils.py
text: @user1:
Shouldn't we be checking that the features that are set by `LanguageModelFeaturizer` are correct, in terms of the mean computed?
We can have a separate test just for `train_utils.align_token_features` which takes as input a random array of features and sub-words as input and computes the aggregated features like the way this test does right now, but for this test I would want to verify if the actual features that are set in the pipeline are correct. Same for ConveRT.

@author:
That should be covered by the other tests in the respective featurizer. The mean in calculated in `train_utils.align_token_features`. It does not matter how the features look like for this method, e.g. if they are coming from ConveRT or BERT.

@user1:
Hmm, then we can have just one test for just `train_utils.align_token_features` right?
In that test we can just feed synthetically created list of tokens? If that's what you had earlier then sorry for the confusion.

@author:
Yeah, can also test on a list of tokens. Will check if `NUMBER_OF_SUB_TOKENS` was correctly set in another test then.