exclude DIET from extractors if no entities trained

Author: tabergma

rasa/nlu/test.py

@@ -37,6 +37,7 @@
 from rasa.nlu.model import Interpreter, Trainer, TrainingData
 from rasa.nlu.components import Component
 from rasa.nlu.tokenizers.tokenizer import Token
+from rasa.utils.tensorflow.constants import ENTITY_RECOGNITION
 
 logger = logging.getLogger(__name__)
 
@@ -1022,12 +1023,17 @@ def get_entity_extractors(interpreter: Interpreter) -> Set[Text]:
 
     Processors are removed since they do not detect the boundaries themselves.
     """
-
     from rasa.nlu.extractors.extractor import EntityExtractor
 
-    extractors = {
-        c.name for c in interpreter.pipeline if isinstance(c, EntityExtractor)
-    }
+    extractors = set()
+    for c in interpreter.pipeline:
+        if isinstance(c, EntityExtractor):
+            if c.name == "DIETClassifier":
+                if c.component_config[ENTITY_RECOGNITION]:
+                    extractors.add(c.name)
+            else:
+                extractors.add(c.name)
+
     return extractors - ENTITY_PROCESSORS
 
 

tests/nlu/test_evaluation.py

@@ -7,6 +7,7 @@
 from _pytest.tmpdir import TempdirFactory
 
 import rasa.utils.io
+from rasa.nlu.classifiers.diet_classifier import DIETClassifier
 from rasa.nlu.extractors.crf_entity_extractor import CRFEntityExtractor
 from rasa.test import compare_nlu_models
 from rasa.nlu.extractors.extractor import EntityExtractor
@@ -50,7 +51,7 @@
 from tests.nlu.conftest import DEFAULT_DATA_PATH
 from rasa.nlu.selectors.response_selector import ResponseSelector
 from rasa.nlu.test import is_response_selector_present
-from rasa.utils.tensorflow.constants import EPOCHS
+from rasa.utils.tensorflow.constants import EPOCHS, ENTITY_RECOGNITION
 
 
 # https://github.com/pytest-dev/pytest-asyncio/issues/68
@@ -510,6 +511,26 @@ def test_response_evaluation_report(tmpdir_factory):
     assert result["predictions"][1] == prediction
 
 
+@pytest.mark.parametrize(
+    "components, expected_extractors",
+    [
+        ([DIETClassifier({ENTITY_RECOGNITION: False})], set()),
+        ([DIETClassifier({ENTITY_RECOGNITION: True})], {"DIETClassifier"}),
+        ([CRFEntityExtractor()], {"CRFEntityExtractor"}),
+        (
+            [SpacyEntityExtractor(), CRFEntityExtractor()],
+            {"SpacyEntityExtractor", "CRFEntityExtractor"},
+        ),
+        ([ResponseSelector()], set()),
+    ],
+)
+def test_get_entity_extractors(components, expected_extractors):
+    mock_interpreter = Interpreter(components, None)
+    extractors = get_entity_extractors(mock_interpreter)
+
+    assert extractors == expected_extractors
+
+
 def test_entity_evaluation_report(tmpdir_factory):
     class EntityExtractorA(EntityExtractor):
 
@@ -653,13 +674,6 @@ def test_evaluate_entities_cv():
     }, "Wrong entity prediction alignment"
 
 
-def test_get_entity_extractors(pretrained_interpreter):
-    assert get_entity_extractors(pretrained_interpreter) == {
-        "SpacyEntityExtractor",
-        "DucklingHTTPExtractor",
-    }
-
-
 def test_remove_pretrained_extractors(pretrained_interpreter):
     target_components_names = ["SpacyNLP"]
     filtered_pipeline = remove_pretrained_extractors(pretrained_interpreter.pipeline)