Repository: RasaHQ/rasa
Pull request: Consider confidence values in 'ensure_consistent_bilou_tagging'
**Proposed changes**:
The method `ensure_consistent_bilou_tagging` now also considers the confidence values of the predicted tags
when updating the BIlOU tags.

closes https://github.com/RasaHQ/rasa/issues/6039

**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:
'ensure_consistent_bilou_tagging' should consider confidence values
**Description of Problem**:
The method `bilou_utils.ensure_consistent_bilou_tagging(tags)` checks if tags are consistently tagged and adapts them if needed. However, confidence values are not modified in this method, although tags might change.

**Overview of the Solution**:
If a tag changed, we should also update the confidence value of that tag.

Hints from issue discussion:

Commit message: add changelog

Patch to review:
diff --git a/changelog/6039.improvment.md b/changelog/6039.improvment.md
new file mode 100644
index 000000000000..a6049d8471b8
--- /dev/null
+++ b/changelog/6039.improvment.md
@@ -0,0 +1,2 @@
+The method `ensure_consistent_bilou_tagging` now also considers the confidence values of the predicted tags
+when updating the BIlOU tags.
\ No newline at end of file
diff --git a/rasa/nlu/classifiers/diet_classifier.py b/rasa/nlu/classifiers/diet_classifier.py
index 1fce5244140f..91320b43d44f 100644
--- a/rasa/nlu/classifiers/diet_classifier.py
+++ b/rasa/nlu/classifiers/diet_classifier.py
@@ -871,7 +871,9 @@ def _entity_label_to_tags(
             tags = [tag_spec.ids_to_tags[p] for p in predictions[0]]
 
             if self.component_config[BILOU_FLAG]:
-                tags = bilou_utils.ensure_consistent_bilou_tagging(tags)
+                tags, confidences = bilou_utils.ensure_consistent_bilou_tagging(
+                    tags, confidences
+                )
 
             predicted_tags[tag_spec.tag_name] = tags
             confidence_values[tag_spec.tag_name] = confidences
diff --git a/rasa/nlu/extractors/crf_entity_extractor.py b/rasa/nlu/extractors/crf_entity_extractor.py
index 6a81cd1d0a07..471fa3661c9b 100644
--- a/rasa/nlu/extractors/crf_entity_extractor.py
+++ b/rasa/nlu/extractors/crf_entity_extractor.py
@@ -286,7 +286,9 @@ def _tag_confidences(
             _tags, _confidences = self._most_likely_tag(predicted_tags)
 
             if self.component_config[BILOU_FLAG]:
-                _tags = bilou_utils.ensure_consistent_bilou_tagging(_tags)
+                _tags, _confidences = bilou_utils.ensure_consistent_bilou_tagging(
+                    _tags, _confidences
+                )
 
             confidences[tag_name] = _confidences
             tags[tag_name] = _tags
diff --git a/rasa/nlu/utils/bilou_utils.py b/rasa/nlu/utils/bilou_utils.py
index 3c4d1b07ddf1..8055af07d8e2 100644
--- a/rasa/nlu/utils/bilou_utils.py
+++ b/rasa/nlu/utils/bilou_utils.py
@@ -1,4 +1,5 @@
 import logging
+from collections import defaultdict
 from typing import List, Tuple, Text, Optional, Dict, Any
 
 from rasa.nlu.tokenizers.tokenizer import Token
@@ -248,20 +249,25 @@ def _add_bilou_tags_to_entities(
                 bilou[end_token_idx] = f"{LAST}{label}"
 
 
-def ensure_consistent_bilou_tagging(predicted_tags: List[Text]) -> List[Text]:
+def ensure_consistent_bilou_tagging(
+    predicted_tags: List[Text], predicted_confidences: List[float]
+) -> Tuple[List[Text], List[float]]:
     """
     Ensure predicted tags follow the BILOU tagging schema.
 
     We assume that starting B- tags are correct. Followed tags that belong to start
-    tag but have a different entity type are updated.
+    tag but have a different entity type are updated considering also the confidence
+    values of those tags.
     For example, B-a I-b L-a is updated to B-a I-a L-a and B-a I-a O is changed to
     B-a L-a.
 
     Args:
         predicted_tags: predicted tags
+        predicted_confidences: predicted confidences
 
     Return:
         List of tags.
+        List of confidences.
     """
 
     for idx, predicted_tag in enumerate(predicted_tags):
@@ -271,6 +277,32 @@ def ensure_consistent_bilou_tagging(predicted_tags: List[Text]) -> List[Text]:
         if prefix == BEGINNING:
             last_idx = _find_bilou_end(idx, predicted_tags)
 
+            relevant_confidences = predicted_confidences[idx : last_idx + 1]
+            relevant_tags = [
+                tag_without_prefix(tag) for tag in predicted_tags[idx : last_idx + 1]
+            ]
+
+            # if not all tags are the same, for example, B-person I-person L-location
+            # we need to check what tag we should use depending on the confidence
+            # values and update the tags and confidences accordingly
+            if not all(relevant_tags[0] == tag for tag in relevant_tags):
+                avg_confidence_per_tag = _avg_confidence_per_tag(
+                    relevant_tags, relevant_confidences
+                )
+
+                # the entity should have the tag with the highest average confidence
+                tag = max(avg_confidence_per_tag, key=avg_confidence_per_tag.get)
+                avg_confidence = avg_confidence_per_tag[tag]
+
+                predicted_confidences = _update_confidences(
+                    predicted_confidences,
+                    predicted_tags,
+                    tag,
+                    avg_confidence,
+                    idx,
+                    last_idx,
+                )
+
             # ensure correct BILOU annotations
             if last_idx == idx:
                 predicted_tags[idx] = f"{UNIT}{tag}"
@@ -283,10 +315,70 @@ def ensure_consistent_bilou_tagging(predicted_tags: List[Text]) -> List[Text]:
                 for i in range(idx + 1, last_idx):
                     predicted_tags[i] = f"{INSIDE}{tag}"
 
-    return predicted_tags
+    return predicted_tags, predicted_confidences
+
+
+def _update_confidences(
+    predicted_confidences: List[float],
+    predicted_tags: List[Text],
+    tag: Text,
+    avg_confidence: float,
+    idx: int,
+    last_idx: int,
+):
+    """Update the confidence values.
+
+    Set the confidence value of a tag to the average confidence value if the predicated
+    tag changed.
+
+    Args:
+        predicted_confidences: The list of predicted confidences.
+        predicted_tags: The list of predicted tags.
+        tag: The tag of the entity.
+        avg_confidence: The average confidence value.
+        idx: The start index of the entity.
+        last_idx: The end index of the entity.
+
+    Returns:
+        The updated list of confidences.
+    """
+    for i in range(idx, last_idx + 1):
+        predicted_confidences[i] = (
+            avg_confidence
+            if tag_without_prefix(predicted_tags[i]) != tag
+            else predicted_confidences[i]
+        )
+    return predicted_confidences
+
+
+def _avg_confidence_per_tag(
+    relevant_tags: List[Text], relevant_confidences: List[float]
+) -> Dict[Text, float]:
+    confidences_per_tag = defaultdict(list)
+
+    for tag, confidence in zip(relevant_tags, relevant_confidences):
+        confidences_per_tag[tag].append(confidence)
+
+    avg_confidence_per_tag = {}
+    for tag, confidences in confidences_per_tag.items():
+        avg_confidence_per_tag[tag] = round(sum(confidences) / len(confidences), 2)
+
+    return avg_confidence_per_tag
 
 
 def _find_bilou_end(start_idx: int, predicted_tags: List[Text]) -> int:
+    """Find the last index of the entity.
+
+    The start index is pointing to a B- tag. The entity is closed as soon as we find
+    a L- tag or a O tag.
+
+    Args:
+        start_idx: The start index of the entity
+        predicted_tags: The list of predicted tags
+
+    Returns:
+        The end index of the entity
+    """
     current_idx = start_idx + 1
     finished = False
     start_tag = tag_without_prefix(predicted_tags[start_idx])
diff --git a/tests/nlu/utils/test_bilou_utils.py b/tests/nlu/utils/test_bilou_utils.py
index 11741bad7a66..3f3b5e7af973 100644
--- a/tests/nlu/utils/test_bilou_utils.py
+++ b/tests/nlu/utils/test_bilou_utils.py
@@ -145,51 +145,93 @@ def test_apply_bilou_schema():
 
 
 @pytest.mark.parametrize(
-    "tags, expected_tags, debug_message",
+    "tags, confidences, expected_tags, expected_confidences, debug_message",
     [
         (
             ["O", "B-person", "I-person", "L-person", "O", "U-person", "O"],
+            [0.99, 0.89, 0.93, 0.99, 0.89, 0.97, 0.87],
             ["O", "B-person", "I-person", "L-person", "O", "U-person", "O"],
+            [0.99, 0.89, 0.93, 0.99, 0.89, 0.97, 0.87],
             None,
         ),
         (
             ["O", "B-person", "B-location", "I-location", "O"],
+            [0.99, 0.89, 0.93, 0.78, 0.89],
             ["O", "U-person", "B-location", "L-location", "O"],
+            [0.99, 0.89, 0.93, 0.78, 0.89],
             "B- tag not closed",
         ),
         (
             ["O", "B-person", "I-location", "L-person"],
+            [0.99, 0.89, 0.77, 0.87],
             ["O", "B-person", "I-person", "L-person"],
+            [0.99, 0.89, 0.88, 0.87],
             "B- tag, L- tag pair encloses multiple entity classes",
         ),
-        (["O", "B-person", "O"], ["O", "U-person", "O"], "B- tag not closed"),
-        (["O", "B-person"], ["O", "U-person"], "B- tag not closed"),
+        (
+            ["O", "B-person", "I-location", "L-location"],
+            [0.99, 0.78, 0.93, 0.96],
+            ["O", "B-location", "I-location", "L-location"],
+            [0.99, 0.95, 0.93, 0.96],
+            "B- tag, L- tag pair encloses multiple entity classes",
+        ),
+        (
+            ["O", "B-person", "I-location", "L-location", "B-person", "L-person"],
+            [0.99, 0.78, 0.93, 0.96, 0.93, 0.96],
+            ["O", "B-location", "I-location", "L-location", "B-person", "L-person"],
+            [0.99, 0.95, 0.93, 0.96, 0.93, 0.96],
+            "B- tag, L- tag pair encloses multiple entity classes",
+        ),
+        (
+            ["O", "B-person", "O"],
+            [0.99, 0.89, 0.87],
+            ["O", "U-person", "O"],
+            [0.99, 0.89, 0.87],
+            "B- tag not closed",
+        ),
+        (
+            ["O", "B-person"],
+            [0.99, 0.89],
+            ["O", "U-person"],
+            [0.99, 0.89],
+            "B- tag not closed",
+        ),
         (
             ["O", "B-person", "I-person"],
+            [0.99, 0.89, 0.87],
             ["O", "B-person", "L-person"],
+            [0.99, 0.89, 0.87],
             "B- tag not closed",
         ),
         (
             ["O", "B-person", "I-location"],
+            [0.99, 0.89, 0.78],
             ["O", "B-person", "L-person"],
+            [0.99, 0.89, 0.89],
             "B- tag not closed",
         ),
         (
             ["O", "B-person", "B-location"],
+            [0.99, 0.89, 0.89],
             ["O", "U-person", "U-location"],
+            [0.99, 0.89, 0.89],
             "B- tag not closed",
         ),
     ],
 )
 def test_check_consistent_bilou_tagging(
     tags: List[Text],
+    confidences: List[float],
     expected_tags: List[Text],
+    expected_confidences: List[float],
     debug_message: Optional[Text],
     caplog: LogCaptureFixture,
 ):
 
     with caplog.at_level(logging.DEBUG):
-        actual_tags = bilou_utils.ensure_consistent_bilou_tagging(tags)
+        actual_tags, actual_confidences = bilou_utils.ensure_consistent_bilou_tagging(
+            tags, confidences
+        )
 
     if debug_message:
         assert len(caplog.records) > 0
@@ -198,3 +240,4 @@ def test_check_consistent_bilou_tagging(
         assert len(caplog.records) == 0
 
     assert actual_tags == expected_tags
+    assert actual_confidences == expected_confidences


Code review comments:
diff_hunk: @@ -271,6 +277,32 @@ def ensure_consistent_bilou_tagging(predicted_tags: List[Text]) -> List[Text]:
         if prefix == BEGINNING:
             last_idx = _find_bilou_end(idx, predicted_tags)
 
+            relevant_confidences = predicted_confidences[idx : last_idx + 1]
+            relevant_tags = [
+                tag_without_prefix(tag) for tag in predicted_tags[idx : last_idx + 1]
+            ]
+
+            # if not all tags are the same, for example, B-person I-person L-location
+            # we need to check what tag we should use depending on the confidence
+            # values and update the tags and confidences accordingly
+            if not all(relevant_tags[0] == tag for tag in relevant_tags):
+                avg_confidence_per_tag = _avg_confidence_per_tag(
+                    relevant_tags, relevant_confidences
+                )
+
+                # the entity should have the tag with the highest average confidence
+                tag = max(avg_confidence_per_tag, key=avg_confidence_per_tag.get)
path: rasa/nlu/utils/bilou_utils.py
text: @author:
@user1 do you think this is enough or should we also take into account how many tags within the entity are predicted as e.g. person or location?

@user1:
I think how many tags is more robust measure. E.g. person 2 times with 0.89 and location 1 time with 0.95. Maybe it is possible to come up with some kind of F1 score between these values

@user1:
also, it'd be great to log this mistakes somehow, so that users can be pointed out where more data is required

@author:
👍 

@author:
updated the metric