Merge pull request #5797 from RasaHQ/empty-role-group

Don't set roles/groups to 'O'
RasaHQ · May 12, 2020 · 79cb399 · 79cb399
2 parents 507d63c + 4abccd6
commit 79cb399
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/changelog/5788.misc.rst b/changelog/5788.misc.rst
@@ -0,0 +1 @@
+Don't set roles and groups to ``O`` (nothing found) when constructing entities from model predictions.
diff --git a/rasa/nlu/extractors/extractor.py b/rasa/nlu/extractors/extractor.py
@@ -507,13 +507,13 @@ def _create_new_entity(
                 ENTITY_ATTRIBUTE_TYPE
             ][idx]
 
-        if ENTITY_ATTRIBUTE_ROLE in tag_names:
+        if ENTITY_ATTRIBUTE_ROLE in tag_names and role_tag != NO_ENTITY_TAG:
             entity[ENTITY_ATTRIBUTE_ROLE] = role_tag
             if confidences is not None:
                 entity[ENTITY_ATTRIBUTE_CONFIDENCE_ROLE] = confidences[
                     ENTITY_ATTRIBUTE_ROLE
                 ][idx]
-        if ENTITY_ATTRIBUTE_GROUP in tag_names:
+        if ENTITY_ATTRIBUTE_GROUP in tag_names and group_tag != NO_ENTITY_TAG:
             entity[ENTITY_ATTRIBUTE_GROUP] = group_tag
             if confidences is not None:
                 entity[ENTITY_ATTRIBUTE_CONFIDENCE_GROUP] = confidences[

diff --git a/tests/nlu/extractors/test_extractor.py b/tests/nlu/extractors/test_extractor.py
@@ -308,6 +308,12 @@ def test_clean_up_entities(
                 },
             ],
         ),
+        (
+            "Amsterdam",
+            {"entity": ["city"], "role": ["O"], "group": ["O"]},
+            None,
+            [{"entity": "city", "start": 0, "end": 9, "value": "Amsterdam"}],
+        ),
     ],
 )
 def test_convert_tags_to_entities(