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

rasa/engine/validation.py

@@ -1,5 +1,6 @@
 import dataclasses
 import inspect
+import re
 import logging
 import sys
 import typing
@@ -30,6 +31,9 @@ from rasa.dialogue_understanding.coexistence.constants import (
     STICKY,
     NON_STICKY,
 )
+from rasa.dialogue_understanding.generator import (
+    LLMBasedCommandGenerator,
+)
 from rasa.dialogue_understanding.patterns.chitchat import FLOW_PATTERN_CHITCHAT
 from rasa.engine.constants import RESERVED_PLACEHOLDERS
 from rasa.engine.exceptions import GraphSchemaValidationException
@@ -51,6 +55,10 @@ from rasa.shared.core.slots import Slot
 from rasa.shared.exceptions import RasaException
 from rasa.shared.nlu.training_data.message import Message
 
+from rasa.dialogue_understanding.coexistence.intent_based_router import (
+    IntentBasedRouter,
+)
+from rasa.dialogue_understanding.coexistence.llm_based_router import LLMBasedRouter
 
 TypeAnnotation = Union[TypeVar, Text, Type, Optional[AvailableEndpoints]]
 
@@ -209,7 +217,7 @@ def _validate_interface_usage(node: SchemaNode) -> None:
         raise GraphSchemaValidationException(
             f"Your model uses a component with class '{node.uses.__name__}'. "
             f"This class does not implement the '{GraphComponent.__name__}' interface "
-            f"and can hence not be run within Rasa Open Source. Please use a different "
+            f"and can hence not be run within Rasa Pro. Please use a different "
             f"component or implement the '{GraphComponent}' interface in class "
             f"'{node.uses.__name__}'. "
             f"See {DOCS_URL_GRAPH_COMPONENTS} for more information."
@@ -503,7 +511,6 @@ def _validate_parent_return_type(
     parent_return_type: TypeAnnotation,
     required_type: TypeAnnotation,
 ) -> None:
-
     if not typing_utils.issubtype(parent_return_type, required_type):
         parent_node_text = ""
         if parent_node:
@@ -606,7 +613,6 @@ def _recursively_check_required_components(
 def validate_flow_component_dependencies(
     flows: FlowsList, model_configuration: GraphModelConfiguration
 ) -> None:
-
     if (pattern_chitchat := flows.flow_by_id(FLOW_PATTERN_CHITCHAT)) is not None:
         _validate_chitchat_dependencies(pattern_chitchat, model_configuration)
 
@@ -637,166 +643,169 @@ def _validate_chitchat_dependencies(
         )
 
 
-def validate_coexistance_routing_setup(
-    domain: Domain, model_configuration: GraphModelConfiguration, flows: FlowsList
+def get_component_index(schema: GraphSchema, component_class: Type) -> Optional[int]:
+    """Extracts the index of a component of the given class in the schema.
+    This function assumes that each component's node name is stored in a way
+    that includes the index as part of the name, formatted as
+    "run_{ComponentName}{Index}", which is how it's created by the recipe.
+    """
+    # the index of the component is at the end of the node name
+    pattern = re.compile(r"\d+$")
+    for node_name, node in schema.nodes.items():
+        if issubclass(node.uses, component_class):
+            match = pattern.search(node_name)
+            if match:
+                index = int(match.group())
+                return index
+    # index is not found or there is no component with the given class
+    return None
+
+
+def get_component_config(
+    schema: GraphSchema, component_class: Type
+) -> Optional[Dict[str, Any]]:
+    """Extracts the config of a component of the given class in the schema."""
+    for node_name, node in schema.nodes.items():
+        if issubclass(node.uses, component_class):
+            return node.config
+    return None
+
+
+def validate_router_exclusivity(schema: GraphSchema) -> None:
+    """Validate that intent-based and llm-based routers are not
+    defined at the same time.
+    """
+    if schema.has_node(IntentBasedRouter) and schema.has_node(LLMBasedRouter):
+        structlogger.error(
+            "validation.coexistance.both_routers_defined",
+            event_info=(
+                "Both LLMBasedRouter and IntentBasedRouter are in the config. "
+                "Please use only one of them."
+            ),
+        )
+        sys.exit(1)
+
+
+def validate_intent_based_router_position(schema: GraphSchema) -> None:
+    """Validate that if intent-based router is defined, it is positioned before
+    the llm command generator.
+    """
+    intent_based_router_pos = get_component_index(schema, IntentBasedRouter)
+    llm_command_generator_pos = get_component_index(schema, LLMBasedCommandGenerator)
+    if (
+        intent_based_router_pos is not None
+        and llm_command_generator_pos is not None
+        and intent_based_router_pos > llm_command_generator_pos
+    ):
+        structlogger.error(
+            "validation.coexistance.wrong_order_of_components",
+            event_info=(
+                "IntentBasedRouter should come before "
+                "a LLMBasedCommandGenerator in the pipeline."
+            ),
+        )
+        sys.exit(1)
+
+
+def validate_that_slots_are_defined_if_router_is_defined(
+    schema: GraphSchema, routing_slots: List[Slot]
 ) -> None:
-    import re
-    from rasa.dialogue_understanding.coexistence.intent_based_router import (
-        IntentBasedRouter,
-    )
-    from rasa.dialogue_understanding.coexistence.llm_based_router import LLMBasedRouter
-    from rasa.dialogue_understanding.generator import LLMCommandGenerator
-
-    def get_component_index(
-        schema: GraphSchema, component_class: Type
-    ) -> Optional[int]:
-        """Extracts the index of a component of the given class in the schema.
-        This function assumes that each component's node name is stored in a way
-        that includes the index as part of the name, formatted as
-        "run_{ComponentName}{Index}", which is how it's created by the recipe.
-        """
-        # the index of the component is at the end of the node name
-        pattern = re.compile(r"\d+$")
-        for node_name, node in schema.nodes.items():
-            if issubclass(node.uses, component_class):
-                match = pattern.search(node_name)
-                if match:
-                    index = int(match.group())
-                    return index
-        # index is not found or there is no component with the given class
-        return None
-
-    def get_component_config(
-        schema: GraphSchema, component_class: Type
-    ) -> Optional[Dict[str, Any]]:
-        """Extracts the config of a component of the given class in the schema."""
-        for node_name, node in schema.nodes.items():
-            if issubclass(node.uses, component_class):
-                return node.config
-        return None
-
-    def validate_router_exclusivity(schema: GraphSchema) -> None:
-        """Validate that intent-based and llm-based routers are not
-        defined at the same time.
-        """
-        if schema.has_node(IntentBasedRouter) and schema.has_node(LLMBasedRouter):
+    # check whether intent-based or llm-based type of router is present
+    for router_type in [IntentBasedRouter, LLMBasedRouter]:
+        router_present = schema.has_node(router_type)
+        slot_has_issue = len(routing_slots) == 0 or routing_slots[0].type_name != "bool"
+        if router_present and slot_has_issue:
             structlogger.error(
-                "validation.coexistance.both_routers_defined",
+                f"validation.coexistance.{ROUTE_TO_CALM_SLOT}_not_in_domain",
                 event_info=(
-                    "Both LLMBasedRouter and IntentBasedRouter are in the config. "
-                    "Please use only one of them."
+                    f"{router_type.__name__} is in the config, but the slot "
+                    f"{ROUTE_TO_CALM_SLOT} is not in the domain or not of "
+                    f"type bool."
                 ),
             )
             sys.exit(1)
 
-    def validate_intent_based_router_position(schema: GraphSchema) -> None:
-        """Validate that if intent-based router is defined, it is positioned before
-        the llm command generator.
-        """
-        intent_based_router_pos = get_component_index(schema, IntentBasedRouter)
-        llm_command_generator_pos = get_component_index(schema, LLMCommandGenerator)
-        if (
-            intent_based_router_pos is not None
-            and llm_command_generator_pos is not None
-            and intent_based_router_pos > llm_command_generator_pos
-        ):
+
+def validate_that_router_is_defined_if_router_slots_are_in_domain(
+    schema: GraphSchema,
+    routing_slots: List[Slot],
+) -> None:
+    defined_router_slots = len(routing_slots) > 0
+    router_present = schema.has_node(IntentBasedRouter) or schema.has_node(
+        LLMBasedRouter
+    )
+    if defined_router_slots and (
+        not router_present or routing_slots[0].type_name != "bool"
+    ):
+        structlogger.error(
+            f"validation.coexistance"
+            f".{ROUTE_TO_CALM_SLOT}_in_domain_with_no_router_defined",
+            event_info=(
+                f"The slot {ROUTE_TO_CALM_SLOT} is in the domain but the "
+                f"LLMBasedRouter or the IntentBasedRouter is not in the config or "
+                f"the type of the slot is not bool."
+            ),
+        )
+        sys.exit(1)
+
+
+def valid_nlu_entry_config(config: Optional[Dict[str, Any]]) -> bool:
+    return (
+        config is not None
+        and NLU_ENTRY in config
+        and isinstance(config[NLU_ENTRY], dict)
+        and STICKY in config[NLU_ENTRY]
+        and NON_STICKY in config[NLU_ENTRY]
+    )
+
+
+def valid_calm_entry_config(config: Optional[Dict[str, Any]]) -> bool:
+    return (
+        config is not None
+        and CALM_ENTRY in config
+        and isinstance(config[CALM_ENTRY], dict)
+        and STICKY in config[CALM_ENTRY]
+    )
+
+
+def validate_configuration(
+    schema: GraphSchema,
+) -> None:
+    """Validate the configuration of the existing coexistence routers."""
+    if schema.has_node(IntentBasedRouter, include_subtypes=False):
+        config = get_component_config(schema, IntentBasedRouter)
+        if not valid_calm_entry_config(config) or not valid_nlu_entry_config(config):
             structlogger.error(
-                "validation.coexistance.wrong_order_of_components",
+                "validation.coexistance.invalid_configuration",
                 event_info=(
-                    "IntentBasedRouter should come before LLMCommandGenerator "
-                    "in the pipeline."
+                    "The configuration of the IntentBasedRouter is invalid. "
+                    "Please check the documentation.",
                 ),
             )
             sys.exit(1)
 
-    def validate_that_slots_are_defined_if_router_is_defined(
-        schema: GraphSchema, routing_slots: List[Slot]
-    ) -> None:
-        # check whether intent-based or llm-based type of router is present
-        for router_type in [IntentBasedRouter, LLMBasedRouter]:
-            router_present = schema.has_node(router_type)
-            slot_has_issue = (
-                len(routing_slots) == 0 or routing_slots[0].type_name != "bool"
-            )
-            if router_present and slot_has_issue:
-                structlogger.error(
-                    f"validation.coexistance.{ROUTE_TO_CALM_SLOT}_not_in_domain",
-                    event_info=(
-                        f"{router_type.__name__} is in the config, but the slot "
-                        f"{ROUTE_TO_CALM_SLOT} is not in the domain or not of "
-                        f"type bool."
-                    ),
-                )
-                sys.exit(1)
-
-    def validate_that_router_is_defined_if_router_slots_are_in_domain(
-        schema: GraphSchema,
-        routing_slots: List[Slot],
-    ) -> None:
-        defined_router_slots = len(routing_slots) > 0
-        router_present = schema.has_node(IntentBasedRouter) or schema.has_node(
-            LLMBasedRouter
-        )
-        if defined_router_slots and (
-            not router_present or routing_slots[0].type_name != "bool"
+    if schema.has_node(LLMBasedRouter, include_subtypes=False):
+        config = get_component_config(schema, LLMBasedRouter)
+        if not valid_calm_entry_config(config) or (
+            config is not None
+            and NLU_ENTRY in config
+            and not valid_nlu_entry_config(config)
         ):
             structlogger.error(
-                f"validation.coexistance"
-                f".{ROUTE_TO_CALM_SLOT}_in_domain_with_no_router_defined",
+                "validation.coexistance.invalid_configuration",
                 event_info=(
-                    f"The slot {ROUTE_TO_CALM_SLOT} is in the domain but the "
-                    f"LLMBasedRouter or the IntentBasedRouter is not in the config or "
-                    f"the type of the slot is not bool."
+                    "The configuration of the LLMBasedRouter is invalid. "
+                    "Please check the documentation.",
                 ),
             )
             sys.exit(1)
 
-    def valid_nlu_entry_config(config: Optional[Dict[str, Any]]) -> bool:
-        return (
-            config is not None
-            and NLU_ENTRY in config
-            and isinstance(config[NLU_ENTRY], dict)
-            and STICKY in config[NLU_ENTRY]
-            and NON_STICKY in config[NLU_ENTRY]
-        )
 
-    def valid_calm_entry_config(config: Optional[Dict[str, Any]]) -> bool:
-        return (
-            config is not None
-            and CALM_ENTRY in config
-            and isinstance(config[CALM_ENTRY], dict)
-            and STICKY in config[CALM_ENTRY]
-        )
-
-    def validate_configuration(
-        schema: GraphSchema,
-    ) -> None:
-        """Validate the configuration of the existing coexistence routers."""
-        if schema.has_node(IntentBasedRouter, include_subtypes=False):
-            config = get_component_config(schema, IntentBasedRouter)
-            if not valid_calm_entry_config(config) or not valid_nlu_entry_config(
-                config
-            ):
-                structlogger.error(
-                    "validation.coexistance.invalid_configuration",
-                    event_info=(
-                        "The configuration of the IntentBasedRouter is invalid. "
-                        "Please check the documentation.",
-                    ),
-                )
-                sys.exit(1)
-
-        if schema.has_node(LLMBasedRouter, include_subtypes=False):
-            config = get_component_config(schema, LLMBasedRouter)
-            if not valid_calm_entry_config(config):
-                structlogger.error(
-                    "validation.coexistance.invalid_configuration",
-                    event_info=(
-                        "The configuration of the LLMBasedRouter is invalid. "
-                        "Please check the documentation.",
-                    ),
-                )
-                sys.exit(1)
+def validate_coexistance_routing_setup(
+    domain: Domain, model_configuration: GraphModelConfiguration, flows: FlowsList
+) -> None:
+    schema = model_configuration.predict_schema
+    routing_slots = [s for s in domain.slots if s.name == ROUTE_TO_CALM_SLOT]
 
     def validate_that_router_or_router_slot_are_defined_if_action_reset_routing_is_used(
         schema: GraphSchema, flows: FlowsList, routing_slots: List[Slot]
@@ -826,9 +835,6 @@ def validate_coexistance_routing_setup(
                 )
             sys.exit(1)
 
-    schema = model_configuration.predict_schema
-    routing_slots = [s for s in domain.slots if s.name == ROUTE_TO_CALM_SLOT]
-
     validate_router_exclusivity(schema)
     validate_intent_based_router_position(schema)
     validate_that_slots_are_defined_if_router_is_defined(schema, routing_slots)
@@ -837,3 +843,31 @@ def validate_coexistance_routing_setup(
     validate_that_router_or_router_slot_are_defined_if_action_reset_routing_is_used(
         schema, flows, routing_slots
     )
+
+
+def validate_command_generator_exclusivity(schema: GraphSchema) -> None:
+    """Validate that multiple command generators are not defined at same time."""
+    from rasa.dialogue_understanding.generator import (
+        LLMBasedCommandGenerator,
+    )
+
+    count = schema.count_nodes_of_a_given_type(
+        LLMBasedCommandGenerator, include_subtypes=True
+    )
+
+    if count > 1:
+        structlogger.error(
+            "validation.command_generator.multiple_command_generator_defined",
+            event_info=(
+                "Multiple LLM based command generators are defined in the config. "
+                "Please use only one LLM based command generator."
+            ),
+        )
+        sys.exit(1)
+
+
+def validate_command_generator_setup(
+    model_configuration: GraphModelConfiguration,
+) -> None:
+    schema = model_configuration.predict_schema
+    validate_command_generator_exclusivity(schema)

rasa/exceptions.py

@@ -20,9 +20,9 @@ class UnsupportedModelVersionError(RasaException):
     def __str__(self) -> Text:
         minimum_version = version.parse(MINIMUM_COMPATIBLE_VERSION)
         return (
-            f"The model version is trained using Rasa Open Source {self.model_version} "
+            f"The model version is trained using Rasa Pro {self.model_version} "
             f"and is not compatible with your current installation "
-            f"which supports models build with Rasa Open Source {minimum_version} "
+            f"which supports models build with Rasa Pro {minimum_version} "
             f"or higher. "
             f"This means that you either need to retrain your model "
             f"or revert back to the Rasa version that trained the model "

rasa/graph_components/validators/default_recipe_validator.py

@@ -203,7 +203,6 @@ class DefaultV1RecipeValidator(GraphComponent):
             )
 
         if training_data.lookup_tables:
-
             if self._component_types.isdisjoint([CRFEntityExtractor, DIETClassifier]):
                 rasa.shared.utils.io.raise_warning(
                     f"You have defined training data consisting of lookup tables, but "
@@ -219,7 +218,6 @@ class DefaultV1RecipeValidator(GraphComponent):
                 )
 
             elif CRFEntityExtractor in self._component_types:
-
                 crf_schema_nodes = [
                     schema_node
                     for schema_node in self._graph_schema.nodes.values()
@@ -295,9 +293,9 @@ class DefaultV1RecipeValidator(GraphComponent):
         Both of these look for the same entities based on the same training data
         leading to ambiguity in the results.
         """
-        extractors_in_configuration: Set[
-            Type[GraphComponent]
-        ] = self._component_types.intersection(TRAINABLE_EXTRACTORS)
+        extractors_in_configuration: Set[Type[GraphComponent]] = (
+            self._component_types.intersection(TRAINABLE_EXTRACTORS)
+        )
         if len(extractors_in_configuration) > 1:
             rasa.shared.utils.io.raise_warning(
                 f"You have defined multiple entity extractors that do the same job "

rasa/hooks.py

@@ -78,7 +78,6 @@ def create_tracker_store(
     domain: "Domain",
     event_broker: Optional["EventBroker"],
 ) -> "TrackerStore":
-
     if isinstance(endpoint_config, EndpointConfig):
         return AuthRetryTrackerStore(
             endpoint_config=endpoint_config, domain=domain, event_broker=event_broker

rasa/model.py

@@ -74,7 +74,7 @@ def get_latest_model(model_path: Text = DEFAULT_MODELS_PATH) -> Optional[Text]:
 
 
 def get_model_for_finetuning(
-    previous_model_file_or_dir: Union[Path, Text]
+    previous_model_file_or_dir: Union[Path, Text],
 ) -> Optional[Path]:
     """Gets validated path for model to finetune.
 

rasa/model_training.py

@@ -309,6 +309,7 @@ async def _train_graph(
     rasa.engine.validation.validate_flow_component_dependencies(
         flows, model_configuration
     )
+    rasa.engine.validation.validate_command_generator_setup(model_configuration)
 
     tempdir_name = rasa.utils.common.get_temp_dir_name()
     # Use `TempDirectoryPath` instead of `tempfile.TemporaryDirectory` as this

rasa/nlu/classifiers/diet_classifier.py

@@ -1,17 +1,18 @@
 from __future__ import annotations
-
 import copy
 import logging
 from collections import defaultdict
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Text, Tuple, Union, TypeVar, Type
+
+from rasa.exceptions import ModelNotFound
+from rasa.nlu.featurizers.featurizer import Featurizer
 
 import numpy as np
 import scipy.sparse
 import tensorflow as tf
 
-from rasa.exceptions import ModelNotFound
-from rasa.nlu.featurizers.featurizer import Featurizer
+from typing import Any, Dict, List, Optional, Text, Tuple, Union, TypeVar, Type
+
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
@@ -19,21 +20,18 @@ from rasa.engine.storage.storage import ModelStorage
 from rasa.nlu.extractors.extractor import EntityExtractorMixin
 from rasa.nlu.classifiers.classifier import IntentClassifier
 import rasa.shared.utils.io
+import rasa.utils.io as io_utils
 import rasa.nlu.utils.bilou_utils as bilou_utils
 from rasa.shared.constants import DIAGNOSTIC_DATA
 from rasa.nlu.extractors.extractor import EntityTagSpec
 from rasa.nlu.classifiers import LABEL_RANKING_LENGTH
 from rasa.utils import train_utils
 from rasa.utils.tensorflow import rasa_layers
-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.model_data import (
     RasaModelData,
     FeatureSignature,
+    FeatureArray,
 )
 from rasa.nlu.constants import TOKENS_NAMES, DEFAULT_TRANSFORMER_SIZE
 from rasa.shared.nlu.constants import (
@@ -120,6 +118,7 @@ LABEL_SUB_KEY = IDS
 
 POSSIBLE_TAGS = [ENTITY_ATTRIBUTE_TYPE, ENTITY_ATTRIBUTE_ROLE, ENTITY_ATTRIBUTE_GROUP]
 
+
 DIETClassifierT = TypeVar("DIETClassifierT", bound="DIETClassifier")
 
 
@@ -511,7 +510,6 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
     def _extract_features(
         self, message: Message, attribute: Text
     ) -> Dict[Text, Union[scipy.sparse.spmatrix, np.ndarray]]:
-
         (
             sparse_sequence_features,
             sparse_sentence_features,
@@ -781,7 +779,6 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
         sparse_feature_sizes: Dict[Text, Dict[Text, List[int]]],
         label_attribute: Optional[Text] = None,
     ) -> Dict[Text, Dict[Text, List[int]]]:
-
         if label_attribute in sparse_feature_sizes:
             del sparse_feature_sizes[label_attribute]
         return sparse_feature_sizes
@@ -1086,24 +1083,18 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
 
             self.model.save(str(tf_model_file))
 
-            # save data example
-            serialize_nested_feature_arrays(
-                self._data_example,
-                model_path / f"{file_name}.data_example.st",
-                model_path / f"{file_name}.data_example_metadata.json",
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.data_example.pkl", self._data_example
             )
-            # save label data
-            serialize_nested_feature_arrays(
-                dict(self._label_data.data) if self._label_data is not None else {},
-                model_path / f"{file_name}.label_data.st",
-                model_path / f"{file_name}.label_data_metadata.json",
-            )
-
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
-                model_path / f"{file_name}.sparse_feature_sizes.json",
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.sparse_feature_sizes.pkl",
                 self._sparse_feature_sizes,
             )
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
+            io_utils.pickle_dump(
+                model_path / f"{file_name}.label_data.pkl",
+                dict(self._label_data.data) if self._label_data is not None else {},
+            )
+            io_utils.json_pickle(
                 model_path / f"{file_name}.index_label_id_mapping.json",
                 self.index_label_id_mapping,
             )
@@ -1192,22 +1183,15 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
     ]:
         file_name = cls.__name__
 
-        # load data example
-        data_example = deserialize_nested_feature_arrays(
-            str(model_path / f"{file_name}.data_example.st"),
-            str(model_path / f"{file_name}.data_example_metadata.json"),
+        data_example = io_utils.pickle_load(
+            model_path / f"{file_name}.data_example.pkl"
         )
-        # load label data
-        loaded_label_data = deserialize_nested_feature_arrays(
-            str(model_path / f"{file_name}.label_data.st"),
-            str(model_path / f"{file_name}.label_data_metadata.json"),
-        )
-        label_data = RasaModelData(data=loaded_label_data)
-
-        sparse_feature_sizes = rasa.shared.utils.io.read_json_file(
-            model_path / f"{file_name}.sparse_feature_sizes.json"
+        label_data = io_utils.pickle_load(model_path / f"{file_name}.label_data.pkl")
+        label_data = RasaModelData(data=label_data)
+        sparse_feature_sizes = io_utils.pickle_load(
+            model_path / f"{file_name}.sparse_feature_sizes.pkl"
         )
-        index_label_id_mapping = rasa.shared.utils.io.read_json_file(
+        index_label_id_mapping = io_utils.json_unpickle(
             model_path / f"{file_name}.index_label_id_mapping.json"
         )
         entity_tag_specs = rasa.shared.utils.io.read_json_file(
@@ -1227,6 +1211,7 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
             for tag_spec in entity_tag_specs
         ]
 
+        # jsonpickle converts dictionary keys to strings
         index_label_id_mapping = {
             int(key): value for key, value in index_label_id_mapping.items()
         }
@@ -1280,7 +1265,6 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
         config: Dict[Text, Any],
         finetune_mode: bool,
     ) -> "RasaModel":
-
         predict_data_example = RasaModelData(
             label_key=model_data_example.label_key,
             data={
@@ -1467,10 +1451,10 @@ class DIET(TransformerRasaModel):
         # everything using a transformer and optionally also do masked language
         # modeling.
         self.text_name = TEXT
-        self._tf_layers[
-            f"sequence_layer.{self.text_name}"
-        ] = rasa_layers.RasaSequenceLayer(
-            self.text_name, self.data_signature[self.text_name], self.config
+        self._tf_layers[f"sequence_layer.{self.text_name}"] = (
+            rasa_layers.RasaSequenceLayer(
+                self.text_name, self.data_signature[self.text_name], self.config
+            )
         )
         if self.config[MASKED_LM]:
             self._prepare_mask_lm_loss(self.text_name)
@@ -1488,10 +1472,10 @@ class DIET(TransformerRasaModel):
                 {SPARSE_INPUT_DROPOUT: False, DENSE_INPUT_DROPOUT: False}
             )
 
-            self._tf_layers[
-                f"feature_combining_layer.{self.label_name}"
-            ] = rasa_layers.RasaFeatureCombiningLayer(
-                self.label_name, self.label_signature[self.label_name], label_config
+            self._tf_layers[f"feature_combining_layer.{self.label_name}"] = (
+                rasa_layers.RasaFeatureCombiningLayer(
+                    self.label_name, self.label_signature[self.label_name], label_config
+                )
             )
 
             self._prepare_ffnn_layer(
@@ -1523,7 +1507,6 @@ class DIET(TransformerRasaModel):
         sequence_feature_lengths: tf.Tensor,
         name: Text,
     ) -> tf.Tensor:
-
         x, _ = self._tf_layers[f"feature_combining_layer.{name}"](
             (sequence_features, sentence_features, sequence_feature_lengths),
             training=self._training,
@@ -1705,7 +1688,6 @@ class DIET(TransformerRasaModel):
         return loss
 
     def _update_label_metrics(self, loss: tf.Tensor, acc: tf.Tensor) -> None:
-
         self.intent_loss.update_state(loss)
         self.intent_acc.update_state(acc)
 
@@ -1864,7 +1846,6 @@ class DIET(TransformerRasaModel):
         combined_sequence_sentence_feature_lengths: tf.Tensor,
         text_transformed: tf.Tensor,
     ) -> Dict[Text, tf.Tensor]:
-
         if self.all_labels_embed is None:
             raise ValueError(
                 "The model was not prepared for prediction. "