Repository: RasaHQ/rasa
Pull request: Rasa test evaluates TED e2e entity extraction
**Proposed changes**:
Fixes: #7578 
Adds functionality to `rasa test` to evaluate the entities TED extracts.

**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/main/changelog) for instructions)
- [x] reformat files using `black` (please check [Readme](https://github.com/RasaHQ/rasa#code-style) for instructions)


Problem statement:
Enable rasa test usage for e2e
Users need `rasa test` to validate and inform their decisions to update their training data. Without it, they cannot have a functioning live project using e2e, which in turns means that we wouldn't be able to get valuable feedback on e2e.

With `e2e` TED can extract entities using data in stories. Currently `rasa test` only runs NLU to compare entities.
We need to update the `rasa test` command so that the entities extracted by TED would be compared with entity labels in stories.

# Enable

**Definition of done (what needs to be true to close this issue)**
- [ ] If TED is configured to extract entities (when using e2e), `rasa test` should compare predictions to golden labels in the stories

Hints from issue discussion:

Commit message: rasa test e2e evluates TED entity extraction

Patch to review:
diff --git a/changelog/7578.improvement.md b/changelog/7578.improvement.md
new file mode 100644
index 000000000000..1ac023a1a660
--- /dev/null
+++ b/changelog/7578.improvement.md
@@ -0,0 +1 @@
+When `TED` is configured to extract entities `rasa test` now evaluates them against the labels in the test stories.  
diff --git a/data/test_e2ebot/config.yml b/data/test_e2ebot/config.yml
new file mode 100644
index 000000000000..4029ca16cc24
--- /dev/null
+++ b/data/test_e2ebot/config.yml
@@ -0,0 +1,16 @@
+language: en
+pipeline:
+  - name: WhitespaceTokenizer
+    intent_tokenization_flag: True
+  - name: RegexFeaturizer
+  - name: LexicalSyntacticFeaturizer
+  - name: CountVectorsFeaturizer
+  - name: CountVectorsFeaturizer
+    analyzer: char_wb
+    min_ngram: 1
+    max_ngram: 4
+  - name: DIETClassifier
+    epochs: 5
+policies:
+- name: TEDPolicy
+  epochs: 5
diff --git a/data/test_e2ebot/data/nlu.yml b/data/test_e2ebot/data/nlu.yml
new file mode 100644
index 000000000000..d103dbff098c
--- /dev/null
+++ b/data/test_e2ebot/data/nlu.yml
@@ -0,0 +1,18 @@
+version: "2.0"
+
+nlu:
+- intent: tell_name
+  examples: |
+    - my name is [joe](name)
+    - my name is [bob](name)
+    - I am [bob](name)
+    - I am [jim](name)
+    - I am called [bill](name)
+
+- intent: tell_mood
+  examples: |
+    - I am feeling [happy](mood)
+    - I am quite [happy](mood)
+    - I feel [sad](mood)
+    - I feel [happy](mood)
+    - I am so [sad](mood)
diff --git a/data/test_e2ebot/data/stories.yml b/data/test_e2ebot/data/stories.yml
new file mode 100644
index 000000000000..adcc69cc686e
--- /dev/null
+++ b/data/test_e2ebot/data/stories.yml
@@ -0,0 +1,16 @@
+version: "2.0"
+
+stories:
+- story: Tell name and mood
+  steps:
+  - user: my name is [joe](name)
+  - bot: hi joe!
+  - user: I am feeling [happy](mood)
+  - bot: I'm glad
+
+- story: Tell name and mood 2
+  steps:
+  - user: my name is [bob](name)
+  - bot: hi bob!
+  - user: I am feeling [sad](mood)
+  - bot: oh dear...
diff --git a/data/test_e2ebot/domain.yml b/data/test_e2ebot/domain.yml
new file mode 100644
index 000000000000..3b97d00b9378
--- /dev/null
+++ b/data/test_e2ebot/domain.yml
@@ -0,0 +1,16 @@
+version: "2.0"
+
+actions:
+ - utter_greet
+
+intents:
+ - tell_name
+ - tell_mood
+
+entities:
+ - name
+ - mood
+
+responses:
+  utter_greet:
+  - text: "hi!"
diff --git a/data/test_e2ebot/tests/test_stories.yml b/data/test_e2ebot/tests/test_stories.yml
new file mode 100644
index 000000000000..3eb9d60aa773
--- /dev/null
+++ b/data/test_e2ebot/tests/test_stories.yml
@@ -0,0 +1,12 @@
+stories:
+- story: Test with two easy entities
+  steps:
+    - user: my name is [joe](name)
+    - bot: hi joe!
+    - user: I am feeling [happy](mood)
+    - bot: I'm glad
+
+- story: Test with a hard entity
+  steps:
+    - user: today I was very [cranky](mood)
+    - bot: I'm sorry!
diff --git a/rasa/core/test.py b/rasa/core/test.py
index 503cb211f4c7..5bd5b4453908 100644
--- a/rasa/core/test.py
+++ b/rasa/core/test.py
@@ -7,7 +7,9 @@
 
 from rasa import telemetry
 from rasa.core.policies.policy import PolicyPrediction
+from rasa.nlu.test import EntityEvaluationResult, evaluate_entities
 from rasa.shared.exceptions import RasaException
+from rasa.shared.nlu.training_data.message import Message
 import rasa.shared.utils.io
 from rasa.core.channels import UserMessage
 from rasa.shared.core.training_data.story_writer.yaml_story_writer import (
@@ -18,6 +20,7 @@
     ENTITY_ATTRIBUTE_TEXT,
     RESPONSE_SELECTOR_DEFAULT_INTENT,
     RESPONSE_SELECTOR_RETRIEVAL_INTENTS,
+    TOKENS_NAMES,
 )
 from rasa.shared.nlu.constants import (
     INTENT,
@@ -32,9 +35,14 @@
     RESPONSE,
     RESPONSE_SELECTOR,
     FULL_RETRIEVAL_INTENT_NAME_KEY,
+    TEXT,
 )
 from rasa.constants import RESULTS_FILE, PERCENTAGE_KEY
-from rasa.shared.core.events import ActionExecuted, UserUttered
+from rasa.shared.core.events import (
+    ActionExecuted,
+    EntitiesAdded,
+    UserUttered,
+)
 from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.nlu.training_data.formats.readerwriter import TrainingDataWriter
 from rasa.shared.importers.importer import TrainingDataImporter
@@ -481,13 +489,38 @@ def emulate_loop_rejection(partial_tracker: DialogueStateTracker) -> None:
     partial_tracker.update(ActionExecutionRejected(rejected_action_name))
 
 
+def _get_processor_entity_evaluation_result(
+    processor: "MessageProcessor",
+    tracker: DialogueStateTracker,
+    prediction: PolicyPrediction,
+):
+    previous_event = tracker.events[-1]
+    if isinstance(previous_event, (UserUttered)):
+        entity_predictions = [
+            entity
+            for prediction_event in prediction.events
+            if isinstance(prediction_event, EntitiesAdded)
+            for entity in prediction_event.entities
+        ]
+        entity_targets = previous_event.entities
+        if entity_targets or entity_predictions:
+            text = previous_event.text
+            parsed_message = processor.interpreter.featurize_message(
+                Message(data={TEXT: text})
+            )
+            tokens = parsed_message.get(TOKENS_NAMES[TEXT])
+            return EntityEvaluationResult(
+                entity_targets, entity_predictions, tokens, text
+            )
+
+
 def _collect_action_executed_predictions(
     processor: "MessageProcessor",
     partial_tracker: DialogueStateTracker,
     event: ActionExecuted,
     fail_on_prediction_errors: bool,
     circuit_breaker_tripped: bool,
-) -> Tuple[EvaluationStore, PolicyPrediction]:
+) -> Tuple[EvaluationStore, PolicyPrediction, Optional[EntityEvaluationResult]]:
     from rasa.core.policies.form_policy import FormPolicy
 
     action_executed_eval_store = EvaluationStore()
@@ -496,6 +529,8 @@ def _collect_action_executed_predictions(
     gold_action_text = event.action_text
     gold = gold_action_name or gold_action_text
 
+    entity_result = None
+
     if circuit_breaker_tripped:
         prediction = PolicyPrediction([], policy_name=None)
         predicted = "circuit breaker tripped"
@@ -503,6 +538,10 @@ def _collect_action_executed_predictions(
         action, prediction = processor.predict_next_action(partial_tracker)
         predicted = action.name()
 
+        entity_result = _get_processor_entity_evaluation_result(
+            processor, partial_tracker, prediction
+        )
+
         if (
             prediction.policy_name
             and predicted != gold
@@ -560,7 +599,7 @@ def _collect_action_executed_predictions(
             )
         )
 
-    return action_executed_eval_store, prediction
+    return action_executed_eval_store, prediction, entity_result
 
 
 def _form_might_have_been_rejected(
@@ -577,7 +616,12 @@ async def _predict_tracker_actions(
     agent: "Agent",
     fail_on_prediction_errors: bool = False,
     use_e2e: bool = False,
-) -> Tuple[EvaluationStore, DialogueStateTracker, List[Dict[Text, Any]]]:
+) -> Tuple[
+    EvaluationStore,
+    DialogueStateTracker,
+    List[Dict[Text, Any]],
+    List[EntityEvaluationResult],
+]:
 
     processor = agent.create_processor()
     tracker_eval_store = EvaluationStore()
@@ -594,19 +638,28 @@ async def _predict_tracker_actions(
     tracker_actions = []
     should_predict_another_action = True
     num_predicted_actions = 0
+    entity_results = []
 
     for event in events[1:]:
         if isinstance(event, ActionExecuted):
             circuit_breaker_tripped = processor.is_action_limit_reached(
                 num_predicted_actions, should_predict_another_action
             )
-            (action_executed_result, prediction) = _collect_action_executed_predictions(
+            (
+                action_executed_result,
+                prediction,
+                entity_result,
+            ) = _collect_action_executed_predictions(
                 processor,
                 partial_tracker,
                 event,
                 fail_on_prediction_errors,
                 circuit_breaker_tripped,
             )
+
+            if entity_result:
+                entity_results.append(entity_result)
+
             tracker_eval_store.merge_store(action_executed_result)
             tracker_actions.append(
                 {
@@ -642,7 +695,7 @@ async def _predict_tracker_actions(
         if isinstance(event, UserUttered):
             num_predicted_actions = 0
 
-    return tracker_eval_store, partial_tracker, tracker_actions
+    return tracker_eval_store, partial_tracker, tracker_actions, entity_results
 
 
 def _in_training_data_fraction(action_list: List[Dict[Text, Any]]) -> float:
@@ -665,7 +718,7 @@ async def _collect_story_predictions(
     agent: "Agent",
     fail_on_prediction_errors: bool = False,
     use_e2e: bool = False,
-) -> Tuple[StoryEvaluation, int]:
+) -> Tuple[StoryEvaluation, int, List[EntityEvaluationResult]]:
     """Test the stories from a file, running them through the stored model."""
     from rasa.test import get_evaluation_metrics
     from tqdm import tqdm
@@ -679,16 +732,20 @@ async def _collect_story_predictions(
     logger.info(f"Evaluating {number_of_stories} stories\nProgress:")
 
     action_list = []
+    entity_results = []
 
     for tracker in tqdm(completed_trackers):
         (
             tracker_results,
             predicted_tracker,
             tracker_actions,
+            tracker_entity_results,
         ) = await _predict_tracker_actions(
             tracker, agent, fail_on_prediction_errors, use_e2e
         )
 
+        entity_results.extend(tracker_entity_results)
+
         story_eval_store.merge_store(tracker_results)
 
         action_list.extend(tracker_actions)
@@ -732,6 +789,7 @@ async def _collect_story_predictions(
             in_training_data_fraction=in_training_data_fraction,
         ),
         number_of_stories,
+        entity_results,
     )
 
 
@@ -781,7 +839,7 @@ async def test(
     generator = await _create_data_generator(stories, agent, max_stories, e2e)
     completed_trackers = generator.generate_story_trackers()
 
-    story_evaluation, _ = await _collect_story_predictions(
+    story_evaluation, _, entity_results = await _collect_story_predictions(
         completed_trackers, agent, fail_on_prediction_errors, e2e
     )
 
@@ -802,11 +860,21 @@ async def test(
             report_filename = os.path.join(out_directory, REPORT_STORIES_FILE)
             rasa.shared.utils.io.dump_obj_as_json_to_file(report_filename, report)
             logger.info(f"Stories report saved to {report_filename}.")
+
         else:
             report, precision, f1, accuracy = get_evaluation_metrics(
                 targets, predictions, output_dict=True
             )
 
+        evaluate_entities(
+            entity_results,
+            {"TEDPolicy"},
+            out_directory,
+            successes,
+            errors,
+            disable_plotting,
+        )
+
     telemetry.track_core_model_test(len(generator.story_graph.story_steps), e2e, agent)
 
     _log_evaluation_table(
@@ -955,7 +1023,8 @@ async def _evaluate_core_model(model: Text, stories_file: Text) -> int:
     agent = Agent.load(model)
     generator = await _create_data_generator(stories_file, agent)
     completed_trackers = generator.generate_story_trackers()
-    story_eval_store, number_of_stories = await _collect_story_predictions(
+
+    story_eval_store, number_of_stories, _ = await _collect_story_predictions(
         completed_trackers, agent
     )
     failed_stories = story_eval_store.failed_stories
diff --git a/tests/conftest.py b/tests/conftest.py
index b17c782f9973..6366759e5da1 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -403,6 +403,24 @@ async def trained_response_selector_bot(trained_async: Callable) -> Path:
     return Path(zipped_model)
 
 
+@pytest.fixture(scope="session")
+async def e2e_bot(trained_async: Callable) -> Path:
+    zipped_model = await trained_async(
+        domain="data/test_e2ebot/domain.yml",
+        config="data/test_e2ebot/config.yml",
+        training_files=[
+            "data/test_e2ebot/data/rules.yml",
+            "data/test_e2ebot/data/stories.yml",
+            "data/test_e2ebot/data/nlu.yml",
+        ],
+    )
+
+    if not zipped_model:
+        raise RasaException("Model training for e2ebot failed.")
+
+    return Path(zipped_model)
+
+
 @pytest.fixture(scope="session")
 async def response_selector_agent(
     trained_response_selector_bot: Optional[Path],
@@ -410,6 +428,11 @@ async def response_selector_agent(
     return Agent.load_local_model(trained_response_selector_bot)
 
 
+@pytest.fixture(scope="session")
+async def e2e_bot_agent(e2e_bot: Optional[Path],) -> Agent:
+    return Agent.load_local_model(e2e_bot)
+
+
 def write_endpoint_config_to_yaml(
     path: Path, data: Dict[Text, Any], endpoints_filename: Text = "endpoints.yml"
 ) -> Path:
diff --git a/tests/core/test_evaluation.py b/tests/core/test_evaluation.py
index e4e0e02a3bd3..f663ad0a4ce3 100644
--- a/tests/core/test_evaluation.py
+++ b/tests/core/test_evaluation.py
@@ -20,7 +20,7 @@
 
 # we need this import to ignore the warning...
 # noinspection PyUnresolvedReferences
-from rasa.nlu.test import run_evaluation
+from rasa.nlu.test import evaluate_entities, run_evaluation
 from rasa.core.agent import Agent
 from tests.core.conftest import (
     DEFAULT_STORIES_FILE,
@@ -60,7 +60,7 @@ async def test_end_to_end_evaluation_script(default_agent: Agent, test_file: Tex
     generator = await _create_data_generator(test_file, default_agent, use_e2e=True)
     completed_trackers = generator.generate_story_trackers()
 
-    story_evaluation, num_stories = await _collect_story_predictions(
+    story_evaluation, num_stories, _ = await _collect_story_predictions(
         completed_trackers, default_agent, use_e2e=True
     )
 
@@ -98,7 +98,7 @@ async def test_end_to_end_evaluation_script_unknown_entity(default_agent: Agent)
     )
     completed_trackers = generator.generate_story_trackers()
 
-    story_evaluation, num_stories = await _collect_story_predictions(
+    story_evaluation, num_stories, _ = await _collect_story_predictions(
         completed_trackers, default_agent, use_e2e=True
     )
 
@@ -114,7 +114,7 @@ async def test_end_to_evaluation_with_forms(form_bot_agent: Agent):
     )
     test_stories = generator.generate_story_trackers()
 
-    story_evaluation, num_stories = await _collect_story_predictions(
+    story_evaluation, num_stories, _ = await _collect_story_predictions(
         test_stories, form_bot_agent, use_e2e=True
     )
 
@@ -153,7 +153,7 @@ async def test_end_to_evaluation_trips_circuit_breaker():
     )
     test_stories = generator.generate_story_trackers()
 
-    story_evaluation, num_stories = await _collect_story_predictions(
+    story_evaluation, num_stories, _ = await _collect_story_predictions(
         test_stories, agent, use_e2e=True
     )
 
@@ -259,7 +259,7 @@ async def test_retrieval_intent(response_selector_agent: Agent, test_file: Text)
     )
     test_stories = generator.generate_story_trackers()
 
-    story_evaluation, num_stories = await _collect_story_predictions(
+    story_evaluation, num_stories, _ = await _collect_story_predictions(
         test_stories, response_selector_agent, use_e2e=True
     )
     # check that test story can either specify base intent or full retrieval intent
@@ -290,3 +290,34 @@ async def test_retrieval_intent_wrong_prediction(
 
     # check if the predicted entry contains full retrieval intent
     assert "# predicted: chitchat/ask_name" in failed_stories
+
+
+async def test_e2e_with_entity_evaluation(e2e_bot_agent: Agent, tmpdir: Path):
+    test_file = "data/test_e2ebot/tests/test_stories.yml"
+
+    await evaluate_stories(
+        stories=test_file,
+        agent=e2e_bot_agent,
+        out_directory=str(tmpdir),
+        max_stories=None,
+        e2e=True,
+    )
+
+    report = rasa.shared.utils.io.read_json_file(tmpdir / "TEDPolicy_report.json")
+    assert report["name"] == {
+        "precision": 1.0,
+        "recall": 1.0,
+        "f1-score": 1.0,
+        "support": 1,
+        "confused_with": {},
+    }
+    assert report["mood"] == {
+        "precision": 1.0,
+        "recall": 0.5,
+        "f1-score": 0.6666666666666666,
+        "support": 2,
+        "confused_with": {},
+    }
+    errors = rasa.shared.utils.io.read_json_file(tmpdir / "TEDPolicy_errors.json")
+    assert len(errors) == 1
+    assert errors[0]["text"] == "today I was very cranky"


Code review comments:
diff_hunk: @@ -481,13 +489,38 @@ def emulate_loop_rejection(partial_tracker: DialogueStateTracker) -> None:
     partial_tracker.update(ActionExecutionRejected(rejected_action_name))
 
 
+def _get_processor_entity_evaluation_result(
path: rasa/core/test.py
text: @user1:
it looks like it only evaluates e2e entities. I think it is better to rename the function to reflect that

diff_hunk: @@ -802,11 +860,21 @@ async def test(
             report_filename = os.path.join(out_directory, REPORT_STORIES_FILE)
             rasa.shared.utils.io.dump_obj_as_json_to_file(report_filename, report)
             logger.info(f"Stories report saved to {report_filename}.")
+
         else:
             report, precision, f1, accuracy = get_evaluation_metrics(
                 targets, predictions, output_dict=True
             )
 
+        evaluate_entities(
+            entity_results,
+            {"TEDPolicy"},
path: rasa/core/test.py
text: @user1:
can we extract the name from entities? I think they contain the name of the extractor

@author:
So find all the unique extractor names in the collection of entities and evaluate them?

@user1:
I meant instead of hard coding `{"TEDPolicy"}`. How do we do it for testing nlu pipeline? We have several extractors there

@author:
For NLU it looks at every component in the pipeline and checks if it is a subclass of `EntityExtractor`. For core we could check the config of every policy for `ENTITY_RECOGNITION=True`? 

@user1:
I think you can just extract source from `EntitiesAdded` event?

@author:
but then we would only evaluate for policies that successfully extracted an entity. 

@user1:
for e2e we need to do exactly that

@user2:
I think what Joe means is that this wouldn't work in case the `TEDPolicy` isn't extracting any entity (and hence not adding an `EntitiesAdded` event)

@author:
yes exactly, and we would want to show that the recall is 0.

@user1:
I see, yeah. Is there a way to know which algorithms are expected to extract these entities? I don't like the idea of hardcoding TED here

@author:
Well we could look for `ENTITY_RECOGNITION` in the config of every policy in the ensemble. But that would only work if that exact config key was always used.

@user1:
maybe for simplicity, in order to unblock us, let's create a constant with a list of policies that can extract e2e entities for now?