Repository: RasaHQ/rasa
Pull request: Fix bug: one word two entity labels
**Proposed changes**:
- closes https://github.com/RasaHQ/rasa/issues/5475

**Status (please check what you already did)**:
- [x] added some tests for the functionality
- [ ] updated the documentation
- [ ] updated the changelog (please check [changelog](https://github.com/RasaHQ/rasa/tree/master/changelog) for instructions)
- [ ] reformat files using `black` (please check [Readme](https://github.com/RasaHQ/rasa#code-style) for instructions)


Problem statement:
One word two entity labels
**Description of the problem**
ConveRT and also other language models we have in our pipeline split words during tokenization into sub-words. `DIETClassifier` assigns different entities to the individual sub-words.

Example:
```
{
    "text": "Aarhus",
    "entities": [
        {
            "start": 0,
            "end": 6,
            "value": "Aarhus",
            "entity": "city"
        }
    ],
    "predicted_entities": [
        {
            "entity": "iata",
            "start": 0,
            "end": 3,
            "extractor": "DIETClassifier",
            "value": "Aar"
        },
        {
            "entity": "city",
            "start": 3,
            "end": 6,
            "extractor": "DIETClassifier",
            "value": "hus"
        }
    ]
```

**Overview of the solution**:
It should not be possible to assign two different entities to one word/token. We should add a sanity check that permits double assignments. We might want to keep the assignment with the higher confidence.

We need to check if this also happens with the `CRFEntityExtractor`.

Hints from issue discussion:

Commit message: fix style issues

Patch to review:
diff --git a/changelog/5475.bugfix.rst b/changelog/5475.bugfix.rst
new file mode 100644
index 000000000000..9f515a0d6421
--- /dev/null
+++ b/changelog/5475.bugfix.rst
@@ -0,0 +1,5 @@
+One word can just have one entity label.
+
+If you are using, for example, ``ConveRTTokenizer`` words can be split into multiple tokens.
+Our entity extractors assign entity labels per token. So, it might happen, that a word, that was split into two tokens,
+got assigned two different entity labels. This is now fixed. One word can just have one entity label at a time.
diff --git a/rasa/nlu/classifiers/diet_classifier.py b/rasa/nlu/classifiers/diet_classifier.py
index 76bfa8a4d39b..c467f9b71b32 100644
--- a/rasa/nlu/classifiers/diet_classifier.py
+++ b/rasa/nlu/classifiers/diet_classifier.py
@@ -743,9 +743,8 @@ def _predict_entities(
 
         return entities
 
-    @staticmethod
     def _convert_tags_to_entities(
-        text: Text, tokens: List[Token], tags: List[Text]
+        self, text: Text, tokens: List[Token], tags: List[Text]
     ) -> List[Dict[Text, Any]]:
         entities = []
         last_tag = NO_ENTITY_TAG
@@ -773,7 +772,7 @@ def _convert_tags_to_entities(
         for entity in entities:
             entity["value"] = text[entity["start"] : entity["end"]]
 
-        return entities
+        return self.clean_up_entities(entities)
 
     def process(self, message: Message, **kwargs: Any) -> None:
         """Return the most likely label and its similarity to the input."""
diff --git a/rasa/nlu/extractors/crf_entity_extractor.py b/rasa/nlu/extractors/crf_entity_extractor.py
index 68b778be70a1..aefa2bea129b 100644
--- a/rasa/nlu/extractors/crf_entity_extractor.py
+++ b/rasa/nlu/extractors/crf_entity_extractor.py
@@ -165,8 +165,9 @@ def extract_entities(self, message: Message) -> List[Dict[Text, Any]]:
         if self.ent_tagger is not None:
             text_data = self._from_text_to_crf(message)
             features = self._sentence_to_features(text_data)
-            ents = self.ent_tagger.predict_marginals_single(features)
-            return self._from_crf_to_json(message, ents)
+            entities = self.ent_tagger.predict_marginals_single(features)
+            entities = self._from_crf_to_json(message, entities)
+            return self.clean_up_entities(entities)
         else:
             return []
 
diff --git a/rasa/nlu/extractors/extractor.py b/rasa/nlu/extractors/extractor.py
index bcdf16ffb366..026ce189d938 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
+from typing import Any, Dict, List, Text, Tuple, Optional
 
 from rasa.nlu.components import Component
 from rasa.nlu.constants import EXTRACTOR, ENTITIES
@@ -21,6 +21,90 @@ def add_processor_name(self, entity: Dict[Text, Any]) -> Dict[Text, Any]:
 
         return entity
 
+    def clean_up_entities(
+        self, entities: List[Dict[Text, Any]], keep: bool = True
+    ) -> List[Dict[Text, Any]]:
+        """
+        Checks if multiple entity labels are assigned to one word.
+
+        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.
+        In such a case keep the entity label with the highest confidence as entity
+        label for that word. If you set 'keep' to 'False', all entity labels for
+        that word will be removed.
+
+        Args:
+            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 list of entities
+        """
+        if len(entities) <= 1:
+            return entities
+
+        entity_indices: List[List[int]] = []
+
+        # get indices of entity labels that belong to one word
+        for idx in range(1, len(entities)):
+            if entities[idx]["start"] == entities[idx - 1]["end"]:
+                if entity_indices and entity_indices[-1][1] == idx - 1:
+                    entity_indices[-1].append(idx)
+                else:
+                    entity_indices.append([idx - 1, idx])
+
+        entities_to_remove = set()
+
+        for indices in entity_indices:
+            if not keep:
+                entities_to_remove.update(indices)
+                continue
+
+            start = entities[indices[0]]["start"]
+            end = entities[indices[-1]]["end"]
+            value = "".join(entities[idx]["value"] for idx in indices)
+            idx = self._get_highest_confidence_idx(entities, indices)
+
+            if idx is None:
+                entities_to_remove.update(indices)
+            else:
+                indices.remove(idx)
+                entities_to_remove.update(indices)
+                entities[idx]["start"] = start
+                entities[idx]["end"] = end
+                entities[idx]["value"] = value
+
+        entities_to_remove = sorted(entities_to_remove, reverse=True)
+
+        for idx in entities_to_remove:
+            entities.remove(entities[idx])
+
+        return entities
+
+    @staticmethod
+    def _get_highest_confidence_idx(
+        entities: List[Dict[Text, Any]], indices: List[int]
+    ) -> Optional[int]:
+        """
+        Args:
+            entities: the full list of entities
+            indices: the indices to consider
+
+        Returns: the idx of the entity label with the highest confidence.
+        """
+        confidences = [
+            entities[idx]["confidence"]
+            for idx in indices
+            if "confidence" in entities[idx]
+        ]
+
+        if len(confidences) != len(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/tests/nlu/extractors/test_extractor.py b/tests/nlu/extractors/test_extractor.py
new file mode 100644
index 000000000000..117bf3709fce
--- /dev/null
+++ b/tests/nlu/extractors/test_extractor.py
@@ -0,0 +1,112 @@
+from typing import Any, Text, Dict, List
+
+import pytest
+
+from rasa.nlu.extractors.extractor import EntityExtractor
+
+
+@pytest.mark.parametrize(
+    "entities, keep, expected_entities",
+    [
+        (
+            [
+                {"entity": "iata", "start": 0, "end": 3, "value": "Aar"},
+                {"entity": "city", "start": 3, "end": 6, "value": "hus"},
+            ],
+            False,
+            [],
+        ),
+        (
+            [
+                {"entity": "iata", "start": 0, "end": 3, "value": "Aar"},
+                {"entity": "city", "start": 3, "end": 6, "value": "hus"},
+            ],
+            True,
+            [],
+        ),
+        (
+            [
+                {"entity": "city", "start": 0, "end": 3, "value": "Aarhus"},
+                {"entity": "type", "start": 4, "end": 9, "value": "city"},
+            ],
+            False,
+            [
+                {"entity": "city", "start": 0, "end": 3, "value": "Aarhus"},
+                {"entity": "type", "start": 4, "end": 9, "value": "city"},
+            ],
+        ),
+        (
+            [
+                {
+                    "entity": "city",
+                    "start": 0,
+                    "end": 3,
+                    "confidence": 0.87,
+                    "value": "Aar",
+                },
+                {
+                    "entity": "iata",
+                    "start": 3,
+                    "end": 6,
+                    "confidence": 0.43,
+                    "value": "hus",
+                },
+            ],
+            True,
+            [
+                {
+                    "entity": "city",
+                    "start": 0,
+                    "end": 6,
+                    "confidence": 0.87,
+                    "value": "Aarhus",
+                }
+            ],
+        ),
+        (
+            [
+                {
+                    "entity": "city",
+                    "start": 0,
+                    "end": 3,
+                    "confidence": 0.87,
+                    "value": "Aar",
+                },
+                {
+                    "entity": "iata",
+                    "start": 3,
+                    "end": 5,
+                    "confidence": 0.43,
+                    "value": "hu",
+                },
+                {
+                    "entity": "city",
+                    "start": 5,
+                    "end": 6,
+                    "confidence": 0.43,
+                    "value": "s",
+                },
+            ],
+            True,
+            [
+                {
+                    "entity": "city",
+                    "start": 0,
+                    "end": 6,
+                    "confidence": 0.87,
+                    "value": "Aarhus",
+                }
+            ],
+        ),
+    ],
+)
+def test_convert_tags_to_entities(
+    entities: List[Dict[Text, Any]],
+    keep: bool,
+    expected_entities: List[Dict[Text, Any]],
+):
+    extractor = EntityExtractor()
+
+    updated_entities = extractor.clean_up_entities(entities, keep)
+
+    assert updated_entities == expected_entities


Code review comments:
diff_hunk: @@ -21,6 +21,90 @@ def add_processor_name(self, entity: Dict[Text, Any]) -> Dict[Text, Any]:
 
         return entity
 
+    def clean_up_entities(
+        self, entities: List[Dict[Text, Any]], keep: bool = True
+    ) -> List[Dict[Text, Any]]:
+        """
+        Checks if multiple entity labels are assigned to one word.
+
+        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.
+        In such a case keep the entity label with the highest confidence as entity
+        label for that word. If you set 'keep' to 'False', all entity labels for
+        that word will be removed.
+
+        Args:
+            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 list of entities
+        """
+        if len(entities) <= 1:
+            return entities
+
+        entity_indices: List[List[int]] = []
+
+        # get indices of entity labels that belong to one word
+        for idx in range(1, len(entities)):
+            if entities[idx]["start"] == entities[idx - 1]["end"]:
+                if entity_indices and entity_indices[-1][1] == idx - 1:
path: rasa/nlu/extractors/extractor.py
text: @user1:
Shouldn't it be `entity_indices[-1][-1] == idx - 1` since you are probably checking for the last index and the last index can be `>1` for words broken into more than 2 tokens?

@author:
good catch, will update the test to catch this scenario