rasa-pro 3.8.18__py3-none-any.whl → 3.9.15__py3-none-any.whl

rasa/core/policies/flows/flow_executor.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 from typing import Any, Dict, Text, List, Optional
 
 from jinja2 import Template
+from rasa.dialogue_understanding.commands import CancelFlowCommand
+from rasa.dialogue_understanding.patterns.cancel import CancelPatternFlowStackFrame
 from structlog.contextvars import (
     bound_contextvars,
 )
@@ -17,6 +19,9 @@ from rasa.core.policies.flows.flow_step_result import (
     FlowStepResult,
     PauseFlowReturnPrediction,
 )
+from rasa.dialogue_understanding.patterns.internal_error import (
+    InternalErrorPatternFlowStackFrame,
+)
 from rasa.dialogue_understanding.patterns.search import SearchPatternFlowStackFrame
 from rasa.dialogue_understanding.stack.dialogue_stack import DialogueStack
 from rasa.dialogue_understanding.stack.frames import (
@@ -42,7 +47,7 @@ from rasa.dialogue_understanding.stack.utils import (
 
 from pypred import Predicate
 
-from rasa.shared.core.constants import ACTION_LISTEN_NAME
+from rasa.shared.core.constants import ACTION_LISTEN_NAME, SlotMappingType
 from rasa.shared.core.events import (
     Event,
     FlowCompleted,
@@ -72,6 +77,7 @@ from rasa.shared.core.flows.flow import (
     FlowStep,
 )
 from rasa.shared.core.flows.steps.collect import SlotRejection
+from rasa.shared.core.slots import Slot
 from rasa.shared.core.trackers import (
     DialogueStateTracker,
 )
@@ -397,7 +403,6 @@ def advance_flows_until_next_action(
     number_of_steps_taken = 0
 
     while isinstance(step_result, ContinueFlowWithNextStep):
-
         number_of_steps_taken += 1
         if number_of_steps_taken > MAX_NUMBER_OF_STEPS:
             raise FlowCircuitBreakerTrippedException(
@@ -467,6 +472,87 @@ def advance_flows_until_next_action(
         return FlowActionPrediction(None, 0.0, events=gathered_events)
 
 
+def validate_collect_step(
+    step: CollectInformationFlowStep,
+    stack: DialogueStack,
+    available_actions: List[str],
+    slots: Dict[Text, Slot],
+) -> bool:
+    """Validate that a collect step can be executed.
+
+    A collect step can be executed if either the `utter_ask` or the `action_ask` is
+    defined in the domain. If neither is defined, the collect step can still be
+    executed if the slot has an initial value defined in the domain, which would cause
+    the step to be skipped."""
+    slot = slots.get(step.collect)
+    slot_has_initial_value_defined = slot and slot.initial_value is not None
+    if (
+        slot_has_initial_value_defined
+        or step.utter in available_actions
+        or step.collect_action in available_actions
+    ):
+        return True
+
+    structlogger.error(
+        "flow.step.run.collect_missing_utter_or_collect_action",
+        slot_name=step.collect,
+    )
+
+    cancel_flow_and_push_internal_error(stack)
+
+    return False
+
+
+def cancel_flow_and_push_internal_error(stack: DialogueStack) -> None:
+    """Cancel the top user flow and push the internal error pattern."""
+    top_frame = stack.top()
+
+    if isinstance(top_frame, BaseFlowStackFrame):
+        # we need to first cancel the top user flow
+        # because we cannot collect one of its slots
+        # and therefore should not proceed with the flow
+        # after triggering pattern_internal_error
+        canceled_frames = CancelFlowCommand.select_canceled_frames(stack)
+        stack.push(
+            CancelPatternFlowStackFrame(
+                canceled_name=top_frame.flow_id,
+                canceled_frames=canceled_frames,
+            )
+        )
+    stack.push(InternalErrorPatternFlowStackFrame())
+
+
+def validate_custom_slot_mappings(
+    step: CollectInformationFlowStep,
+    stack: DialogueStack,
+    tracker: DialogueStateTracker,
+    available_actions: List[str],
+) -> bool:
+    """Validate a slot with custom mappings.
+
+    If invalid, trigger pattern_internal_error and return False.
+    """
+    slot = tracker.slots.get(step.collect, None)
+    slot_mappings = slot.mappings if slot else []
+    for mapping in slot_mappings:
+        if (
+            mapping.get("type") == SlotMappingType.CUSTOM.value
+            and mapping.get("action") is None
+        ):
+            # this is a slot that must be filled by a custom action
+            # check if collect_action exists
+            if step.collect_action not in available_actions:
+                structlogger.error(
+                    "flow.step.run.collect_action_not_found_for_custom_slot_mapping",
+                    action=step.collect_action,
+                    collect=step.collect,
+                )
+                cancel_flow_and_push_internal_error(stack)
+                return False
+
+    return True
+
+
 def run_step(
     step: FlowStep,
     flow: Flow,
@@ -500,6 +586,22 @@ def run_step(
         initial_events.append(FlowStarted(flow.id))
 
     if isinstance(step, CollectInformationFlowStep):
+        is_step_valid = validate_collect_step(
+            step, stack, available_actions, tracker.slots
+        )
+        if not is_step_valid:
+            # if we return any other FlowStepResult, the assistant will stay silent
+            # instead of triggering the internal error pattern
+            return ContinueFlowWithNextStep(events=initial_events)
+
+        is_mapping_valid = validate_custom_slot_mappings(
+            step, stack, tracker, available_actions
+        )
+        if not is_mapping_valid:
+            # if we return any other FlowStepResult, the assistant will stay silent
+            # instead of triggering the internal error pattern
+            return ContinueFlowWithNextStep(events=initial_events)
+
         structlogger.debug("flow.step.run.collect")
         trigger_pattern_ask_collect_information(
             step.collect, stack, step.rejections, step.utter, step.collect_action

rasa/core/policies/intentless_policy.py

@@ -3,7 +3,6 @@ import math
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING, Text, Tuple
 
-import rasa.shared.utils.io
 import structlog
 import tiktoken
 from jinja2 import Template
@@ -11,6 +10,7 @@ from langchain.docstore.document import Document
 from langchain.schema.embeddings import Embeddings
 from langchain.vectorstores import FAISS
 
+import rasa.shared.utils.io
 from rasa import telemetry
 from rasa.core.constants import (
     CHAT_POLICY_PRIORITY,
@@ -56,7 +56,6 @@ from rasa.shared.utils.llm import (
     sanitize_message_for_prompt,
     tracker_as_readable_transcript,
 )
-
 from rasa.utils.ml_utils import (
     extract_ai_response_examples,
     extract_participant_messages_from_transcript,
@@ -65,7 +64,6 @@ from rasa.utils.ml_utils import (
     persist_faiss_vector_store,
     response_for_template,
 )
-
 from rasa.utils.log_utils import log_llm
 
 if TYPE_CHECKING:
@@ -543,7 +541,9 @@ class IntentlessPolicy(Policy):
         Returns:
              The prediction.
         """
-        if not self.supports_current_stack_frame(tracker):
+        if not self.supports_current_stack_frame(
+            tracker
+        ) or self.should_abstain_in_coexistence(tracker, True):
             return self._prediction(self._default_predictions(domain))
 
         if tracker.has_bot_message_after_latest_user_message():
@@ -670,7 +670,7 @@ class IntentlessPolicy(Policy):
         if tracker.latest_message.text.startswith("/"):
             # we don't want to generate a response if the user is trying to
             # execute a "command" - this should be handled by the regex
-            # intent classifier in rasa open source.
+            # intent classifier in rasa pro.
             structlogger.debug("intentless_policy.prediction.skip_slash")
             return None, 0.0
 
@@ -863,7 +863,7 @@ class IntentlessPolicy(Policy):
         """
         result = self._default_predictions(domain)
         if action_name:
-            result[domain.index_for_action(action_name)] = score  # type: ignore[assignment]  # noqa: E501
+            result[domain.index_for_action(action_name)] = score  # type: ignore[assignment]
         return result
 
     @classmethod
@@ -892,9 +892,7 @@ class IntentlessPolicy(Policy):
                 #  normalized. unfortunatley langchain doesn't persist / load
                 #  this parameter.
                 if responses_docsearch:
-                    responses_docsearch._normalize_L2 = (
-                        True  # pylint: disable=protected-access
-                    )
+                    responses_docsearch._normalize_L2 = True  # pylint: disable=protected-access
                 prompt_template = rasa.shared.utils.io.read_file(
                     path / INTENTLESS_PROMPT_TEMPLATE_FILE_NAME
                 )

rasa/core/policies/memoization.py

@@ -419,9 +419,9 @@ class AugmentedMemoizationPolicy(MemoizationPolicy):
         logger.debug("Launch DeLorean...")
 
         # Truncate the tracker based on `max_history`
-        truncated_tracker: Optional[
-            DialogueStateTracker
-        ] = _trim_tracker_by_max_history(tracker, self.config[POLICY_MAX_HISTORY])
+        truncated_tracker: Optional[DialogueStateTracker] = (
+            _trim_tracker_by_max_history(tracker, self.config[POLICY_MAX_HISTORY])
+        )
         truncated_tracker = self._strip_leading_events_until_action_executed(
             truncated_tracker
         )

rasa/core/policies/policy.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
+
 import abc
 import copy
 import logging
 from enum import Enum
 from pathlib import Path
-
-from rasa.shared.constants import ROUTE_TO_CALM_SLOT
-from rasa.shared.core.events import Event
 from typing import (
     Any,
     List,
@@ -21,6 +19,8 @@ from typing import (
 
 import numpy as np
 
+from rasa.shared.constants import ROUTE_TO_CALM_SLOT
+from rasa.shared.core.events import Event
 from rasa.engine.graph import GraphComponent, ExecutionContext
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
@@ -40,14 +40,12 @@ from rasa.core.constants import (
 from rasa.shared.core.constants import USER, SLOTS, PREVIOUS_ACTION, ACTIVE_LOOP
 import rasa.shared.utils.common
 
-
 if TYPE_CHECKING:
     from rasa.shared.nlu.training_data.features import Features
     from rasa.core.featurizers.tracker_featurizers import TrackerFeaturizer
     from rasa.core.featurizers.tracker_featurizers import MaxHistoryTrackerFeaturizer
     from rasa.dialogue_understanding.stack.frames import DialogueStackFrame
 
-
 logger = logging.getLogger(__name__)
 
 TrackerListTypeVar = TypeVar(
@@ -137,10 +135,20 @@ class Policy(GraphComponent):
     def should_abstain_in_coexistence(
         self, tracker: DialogueStateTracker, is_calm_policy: bool
     ) -> bool:
-        """Whether a policy should abstain making predictions in coexistence."""
+        """Whether a policy should abstain making predictions in coexistence.
+
+        A calm policy should run when the routing slot is set to True.
+        A nlu-based policy should run when the routing slot is set to False or None.
+        """
+        if is_calm_policy:
+            return tracker.has_coexistence_routing_slot and (
+                tracker.get_slot(ROUTE_TO_CALM_SLOT) is False
+                or tracker.get_slot(ROUTE_TO_CALM_SLOT) is None
+            )
+
         return (
             tracker.has_coexistence_routing_slot
-            and tracker.get_slot(ROUTE_TO_CALM_SLOT) != is_calm_policy
+            and tracker.get_slot(ROUTE_TO_CALM_SLOT) is True
         )
 
     def __init__(
@@ -299,8 +307,9 @@ class Policy(GraphComponent):
         max_training_samples = kwargs.get("max_training_samples")
         if max_training_samples is not None:
             logger.debug(
-                "Limit training data to {} training samples."
-                "".format(max_training_samples)
+                "Limit training data to {} training samples.".format(
+                    max_training_samples
+                )
             )
             state_features = state_features[:max_training_samples]
             label_ids = label_ids[:max_training_samples]

rasa/core/policies/rule_policy.py

@@ -60,7 +60,7 @@ logger = logging.getLogger(__name__)
 structlogger = structlog.get_logger()
 
 
-# These are Rasa Open Source default actions and overrule everything at any time.
+# These are Rasa Pro default actions and overrule everything at any time.
 DEFAULT_ACTION_MAPPINGS = {
     USER_INTENT_RESTART: ACTION_RESTART_NAME,
     USER_INTENT_BACK: ACTION_BACK_NAME,
@@ -271,15 +271,13 @@ class RulePolicy(MemoizationPolicy):
             if (
                 # loop is predicted after action_listen in unhappy path,
                 # therefore no validation is needed
-                is_prev_action_listen_in_state(states[-1])
-                and action == active_loop
+                is_prev_action_listen_in_state(states[-1]) and action == active_loop
             ):
                 lookup[feature_key] = LOOP_WAS_INTERRUPTED
             elif (
                 # some action other than active_loop is predicted in unhappy path,
                 # therefore active_loop shouldn't be predicted by the rule
-                not is_prev_action_listen_in_state(states[-1])
-                and action != active_loop
+                not is_prev_action_listen_in_state(states[-1]) and action != active_loop
             ):
                 lookup[feature_key] = DO_NOT_PREDICT_LOOP_ACTION
         return lookup
@@ -777,10 +775,10 @@ class RulePolicy(MemoizationPolicy):
         trackers_as_actions = rule_trackers_as_actions + story_trackers_as_actions
 
         # negative rules are not anti-rules, they are auxiliary to actual rules
-        self.lookup[
-            RULES_FOR_LOOP_UNHAPPY_PATH
-        ] = self._create_loop_unhappy_lookup_from_states(
-            trackers_as_states, trackers_as_actions
+        self.lookup[RULES_FOR_LOOP_UNHAPPY_PATH] = (
+            self._create_loop_unhappy_lookup_from_states(
+                trackers_as_states, trackers_as_actions
+            )
         )
 
     def train(
@@ -955,7 +953,6 @@ class RulePolicy(MemoizationPolicy):
     def _find_action_from_loop_happy_path(
         tracker: DialogueStateTracker,
     ) -> Tuple[Optional[Text], Optional[Text]]:
-
         active_loop_name = tracker.active_loop_name
         if active_loop_name is None:
             return None, None
@@ -1132,7 +1129,7 @@ class RulePolicy(MemoizationPolicy):
             tracker, domain, use_text_for_last_user_input=True
         )
 
-        # Rasa Open Source default actions overrule anything. If users want to achieve
+        # Rasa Pro default actions overrule anything. If users want to achieve
         # the same, they need to write a rule or make sure that their loop rejects
         # accordingly.
         (

rasa/core/policies/ted_policy.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
-
 import logging
+
+from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from pathlib import Path
 from collections import defaultdict
 import contextlib
-from typing import Any, List, Optional, Text, Dict, Tuple, Union, Type
 
 import numpy as np
 import tensorflow as tf
+from typing import Any, List, Optional, Text, Dict, Tuple, Union, Type
 
-from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.graph import ExecutionContext
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
@@ -49,22 +49,18 @@ from rasa.shared.core.generator import TrackerWithCachedStates
 from rasa.shared.core.events import EntitiesAdded, Event
 from rasa.shared.core.domain import Domain
 from rasa.shared.nlu.training_data.message import Message
-from rasa.shared.nlu.training_data.features import (
-    Features,
-    save_features,
-    load_features,
-)
+from rasa.shared.nlu.training_data.features import Features
 import rasa.shared.utils.io
 import rasa.utils.io
 from rasa.utils import train_utils
-from rasa.utils.tensorflow.feature_array import (
-    FeatureArray,
-    serialize_nested_feature_arrays,
-    deserialize_nested_feature_arrays,
-)
 from rasa.utils.tensorflow.models import RasaModel, TransformerRasaModel
 from rasa.utils.tensorflow import rasa_layers
-from rasa.utils.tensorflow.model_data import RasaModelData, FeatureSignature, Data
+from rasa.utils.tensorflow.model_data import (
+    RasaModelData,
+    FeatureSignature,
+    FeatureArray,
+    Data,
+)
 from rasa.utils.tensorflow.model_data_utils import convert_to_data_format
 from rasa.utils.tensorflow.constants import (
     LABEL,
@@ -472,7 +468,7 @@ class TEDPolicy(Policy):
 
     @staticmethod
     def _should_extract_entities(
-        entity_tags: List[List[Dict[Text, List[Features]]]]
+        entity_tags: List[List[Dict[Text, List[Features]]]],
     ) -> bool:
         for turns_tags in entity_tags:
             for turn_tags in turns_tags:
@@ -965,32 +961,22 @@ class TEDPolicy(Policy):
             model_path: Path where model is to be persisted
         """
         model_filename = self._metadata_filename()
-        rasa.shared.utils.io.dump_obj_as_json_to_file(
-            model_path / f"{model_filename}.priority.json", self.priority
+        rasa.utils.io.json_pickle(
+            model_path / f"{model_filename}.priority.pkl", self.priority
         )
-        rasa.shared.utils.io.dump_obj_as_json_to_file(
-            model_path / f"{model_filename}.meta.json", self.config
-        )
-        # save data example
-        serialize_nested_feature_arrays(
-            self.data_example,
-            str(model_path / f"{model_filename}.data_example.st"),
-            str(model_path / f"{model_filename}.data_example_metadata.json"),
+        rasa.utils.io.pickle_dump(
+            model_path / f"{model_filename}.meta.pkl", self.config
         )
-        # save label data
-        serialize_nested_feature_arrays(
-            dict(self._label_data.data) if self._label_data is not None else {},
-            str(model_path / f"{model_filename}.label_data.st"),
-            str(model_path / f"{model_filename}.label_data_metadata.json"),
+        rasa.utils.io.pickle_dump(
+            model_path / f"{model_filename}.data_example.pkl", self.data_example
         )
-        # save fake features
-        metadata = save_features(
-            self.fake_features, str(model_path / f"{model_filename}.fake_features.st")
+        rasa.utils.io.pickle_dump(
+            model_path / f"{model_filename}.fake_features.pkl", self.fake_features
         )
-        rasa.shared.utils.io.dump_obj_as_json_to_file(
-            model_path / f"{model_filename}.fake_features_metadata.json", metadata
+        rasa.utils.io.pickle_dump(
+            model_path / f"{model_filename}.label_data.pkl",
+            dict(self._label_data.data) if self._label_data is not None else {},
         )
-
         entity_tag_specs = (
             [tag_spec._asdict() for tag_spec in self._entity_tag_specs]
             if self._entity_tag_specs
@@ -1008,29 +994,18 @@ class TEDPolicy(Policy):
             model_path: Path where model is to be persisted.
         """
         tf_model_file = model_path / f"{cls._metadata_filename()}.tf_model"
-
-        # load data example
-        loaded_data = deserialize_nested_feature_arrays(
-            str(model_path / f"{cls._metadata_filename()}.data_example.st"),
-            str(model_path / f"{cls._metadata_filename()}.data_example_metadata.json"),
+        loaded_data = rasa.utils.io.pickle_load(
+            model_path / f"{cls._metadata_filename()}.data_example.pkl"
         )
-        # load label data
-        loaded_label_data = deserialize_nested_feature_arrays(
-            str(model_path / f"{cls._metadata_filename()}.label_data.st"),
-            str(model_path / f"{cls._metadata_filename()}.label_data_metadata.json"),
+        label_data = rasa.utils.io.pickle_load(
+            model_path / f"{cls._metadata_filename()}.label_data.pkl"
         )
-        label_data = RasaModelData(data=loaded_label_data)
-
-        # load fake features
-        metadata = rasa.shared.utils.io.read_json_file(
-            model_path / f"{cls._metadata_filename()}.fake_features_metadata.json"
-        )
-        fake_features = load_features(
-            str(model_path / f"{cls._metadata_filename()}.fake_features.st"), metadata
+        fake_features = rasa.utils.io.pickle_load(
+            model_path / f"{cls._metadata_filename()}.fake_features.pkl"
         )
-
-        priority = rasa.shared.utils.io.read_json_file(
-            model_path / f"{cls._metadata_filename()}.priority.json"
+        label_data = RasaModelData(data=label_data)
+        priority = rasa.utils.io.json_unpickle(
+            model_path / f"{cls._metadata_filename()}.priority.pkl"
         )
         entity_tag_specs = rasa.shared.utils.io.read_json_file(
             model_path / f"{cls._metadata_filename()}.entity_tag_specs.json"
@@ -1048,8 +1023,8 @@ class TEDPolicy(Policy):
             )
             for tag_spec in entity_tag_specs
         ]
-        model_config = rasa.shared.utils.io.read_json_file(
-            model_path / f"{cls._metadata_filename()}.meta.json"
+        model_config = rasa.utils.io.pickle_load(
+            model_path / f"{cls._metadata_filename()}.meta.pkl"
         )
 
         return {
@@ -1095,7 +1070,7 @@ class TEDPolicy(Policy):
     ) -> TEDPolicy:
         featurizer = TrackerFeaturizer.load(model_path)
 
-        if not (model_path / f"{cls._metadata_filename()}.data_example.st").is_file():
+        if not (model_path / f"{cls._metadata_filename()}.data_example.pkl").is_file():
             return cls(
                 config,
                 model_storage,
@@ -1117,7 +1092,7 @@ class TEDPolicy(Policy):
 
         model = None
 
-        with (contextlib.nullcontext() if config["use_gpu"] else tf.device("/cpu:0")):
+        with contextlib.nullcontext() if config["use_gpu"] else tf.device("/cpu:0"):
             model = cls._load_tf_model(
                 model_utilities,
                 model_data_example,
@@ -1291,19 +1266,19 @@ class TED(TransformerRasaModel):
             )
             self._prepare_encoding_layers(name)
 
-        self._tf_layers[
-            f"transformer.{DIALOGUE}"
-        ] = rasa_layers.prepare_transformer_layer(
-            attribute_name=DIALOGUE,
-            config=self.config,
-            num_layers=self.config[NUM_TRANSFORMER_LAYERS][DIALOGUE],
-            units=self.config[TRANSFORMER_SIZE][DIALOGUE],
-            drop_rate=self.config[DROP_RATE_DIALOGUE],
-            # use bidirectional transformer, because
-            # we will invert dialogue sequence so that the last turn is located
-            # at the first position and would always have
-            # exactly the same positional encoding
-            unidirectional=not self.max_history_featurizer_is_used,
+        self._tf_layers[f"transformer.{DIALOGUE}"] = (
+            rasa_layers.prepare_transformer_layer(
+                attribute_name=DIALOGUE,
+                config=self.config,
+                num_layers=self.config[NUM_TRANSFORMER_LAYERS][DIALOGUE],
+                units=self.config[TRANSFORMER_SIZE][DIALOGUE],
+                drop_rate=self.config[DROP_RATE_DIALOGUE],
+                # use bidirectional transformer, because
+                # we will invert dialogue sequence so that the last turn is located
+                # at the first position and would always have
+                # exactly the same positional encoding
+                unidirectional=not self.max_history_featurizer_is_used,
+            )
         )
 
         self._prepare_label_classification_layers(DIALOGUE)
@@ -1333,23 +1308,23 @@ class TED(TransformerRasaModel):
         # Attributes with sequence-level features also have sentence-level features,
         # all these need to be combined and further processed.
         if attribute_name in SEQUENCE_FEATURES_TO_ENCODE:
-            self._tf_layers[
-                f"sequence_layer.{attribute_name}"
-            ] = rasa_layers.RasaSequenceLayer(
-                attribute_name, attribute_signature, config_to_use
+            self._tf_layers[f"sequence_layer.{attribute_name}"] = (
+                rasa_layers.RasaSequenceLayer(
+                    attribute_name, attribute_signature, config_to_use
+                )
             )
         # Attributes without sequence-level features require some actual feature
         # processing only if they have sentence-level features. Attributes with no
         # sequence- and sentence-level features (dialogue, entity_tags, label) are
         # skipped here.
         elif SENTENCE in attribute_signature:
-            self._tf_layers[
-                f"sparse_dense_concat_layer.{attribute_name}"
-            ] = rasa_layers.ConcatenateSparseDenseFeatures(
-                attribute=attribute_name,
-                feature_type=SENTENCE,
-                feature_type_signature=attribute_signature[SENTENCE],
-                config=config_to_use,
+            self._tf_layers[f"sparse_dense_concat_layer.{attribute_name}"] = (
+                rasa_layers.ConcatenateSparseDenseFeatures(
+                    attribute=attribute_name,
+                    feature_type=SENTENCE,
+                    feature_type_signature=attribute_signature[SENTENCE],
+                    config=config_to_use,
+                )
             )
 
     def _prepare_encoding_layers(self, name: Text) -> None:
@@ -1385,7 +1360,7 @@ class TED(TransformerRasaModel):
 
     @staticmethod
     def _compute_dialogue_indices(
-        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]]
+        tf_batch_data: Dict[Text, Dict[Text, List[tf.Tensor]]],
     ) -> None:
         dialogue_lengths = tf.cast(tf_batch_data[DIALOGUE][LENGTH][0], dtype=tf.int32)
         # wrap in a list, because that's the structure of tf_batch_data
@@ -1424,7 +1399,7 @@ class TED(TransformerRasaModel):
 
     @staticmethod
     def _collect_label_attribute_encodings(
-        all_labels_encoded: Dict[Text, tf.Tensor]
+        all_labels_encoded: Dict[Text, tf.Tensor],
     ) -> tf.Tensor:
         # Initialize with at least one attribute first
         # so that the subsequent TF ops are simplified.
@@ -1953,7 +1928,6 @@ class TED(TransformerRasaModel):
         text_output: tf.Tensor,
         text_sequence_lengths: tf.Tensor,
     ) -> tf.Tensor:
-
         text_transformed, text_mask, text_sequence_lengths = self._reshape_for_entities(
             tf_batch_data,
             dialogue_transformer_output,
@@ -2156,7 +2130,6 @@ class TED(TransformerRasaModel):
         text_output: tf.Tensor,
         text_sequence_lengths: tf.Tensor,
     ) -> Tuple[tf.Tensor, tf.Tensor]:
-
         text_transformed, _, text_sequence_lengths = self._reshape_for_entities(
             tf_batch_data,
             dialogue_transformer_output,