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

README.md

@@ -236,6 +236,39 @@ To check the types execute
 make types
 ```
 
+### Backporting
+
+In order to port changes to `main` and across release branches, we use the `backport` workflow located at
+the `.github/workflows/backport.yml` path.
+This workflow is triggered by the `backport-to-<release-branch>` label applied to a PR, for example `backport-to-3.8.x`.
+Current available target branches are `main` and maintained release branches.
+
+When a PR gets labelled `backport-to-<release-branch>`, a PR is opened by the `backport-github-action` as soon as the
+source PR gets closed (by merging). If you want to close the PR without merging changes, make sure to remove the `backport-to-<release-branch>` label.
+
+The PR author which the action assigns to the backporting PR has to resolve any conflicts before approving and merging.
+Release PRs should also be labelled with `backport-to-main` to backport the `CHANGELOG.md` updates to `main`.
+Backporting version updates should be accepted to the `main` branch from the latest release branch only.
+
+Here are some guidelines to follow when backporting changes and resolving conflicts:
+
+a) for conflicts in `version.py`: accept only the version from the latest release branch. Do not merge version changes
+from earlier release branches into `main` because this could cause issues when trying to make the next minor release.
+
+b) for conflicts in `pyproject.toml`: if related to the `rasa-pro` version, accept only the latest release branch; 
+if related to other dependencies, accept `main` or whichever is the higher upgrade (main usually has the updated 
+dependencies because we only do housekeeping on `main`, apart from vulnerability updates). Be mindful of dependencies that
+are removed from `main` but still exist in former release branches (for example `langchain`).
+
+c) for conflicts in `poetry.lock`: accept changes which were already present on the target branch, then run
+`poetry lock --no-update` so that the lock file contains your changes from `pyproject.toml` too.
+
+d) for conflicts in `CHANGELOG.md`: Manually place the changelog in their allocated section (e.g. 3.8.10 will go under the 
+3.8 section with the other releases, rather than go at the top of the file)
+
+If the backporting workflow fails, you are encouraged to cherry-pick the commits manually and create a PR to
+the target branch. Alternatively, you can install the backporting CLI tool as described [here](https://github.com/sorenlouv/backport?tab=readme-ov-file#install).
+
 ## Releases
 Rasa has implemented robust policies governing version naming, as well as release pace for major, minor, and patch releases.
 
@@ -318,9 +351,12 @@ Releasing a new version is quite simple, as the packages are build and distribut
 9. If however an error occurs in the build, then we should see a failure message automatically posted in the company's Slack (`dev-tribe` channel) like this [one](https://rasa-hq.slack.com/archives/C01M5TAHDHA/p1701444735622919)
    (In this case do the following checks):
     - Check the workflows in [Github Actions](https://github.com/RasaHQ/rasa-private/actions) and make sure that the merged PR of the current release is completed successfully. To easily find your PR you can use the filters `event: push` and `branch: <version number>` (example on release 2.4 you can see [here](https://github.com/RasaHQ/rasa/actions/runs/643344876))
-    - If the workflow is not completed, then try to re run the workflow in case that solves the problem
+    - If the workflow is not completed, then try to re-run the workflow in case that solves the problem
     - If the problem persists, check also the log files and try to find the root cause of the issue
     - If you still cannot resolve the error, contact the infrastructure team by providing any helpful information from your investigation
+10. If the release is successful, add the newly created release branch to the backporting configuration in the `.backportrc.json` file to
+the `targetBranchesChoices` list. This is necessary for the backporting workflow to work correctly with new release branches.
+
 
 ### Cutting a Patch release
 

rasa/constants.py

@@ -18,7 +18,7 @@ CONFIG_TELEMETRY_ID = "rasa_user_id"
 CONFIG_TELEMETRY_ENABLED = "enabled"
 CONFIG_TELEMETRY_DATE = "date"
 
-MINIMUM_COMPATIBLE_VERSION = "3.7.0"
+MINIMUM_COMPATIBLE_VERSION = "3.9.16"
 
 GLOBAL_USER_CONFIG_PATH = os.path.expanduser("~/.config/rasa/global.yml")
 

rasa/core/featurizers/single_state_featurizer.py

@@ -1,7 +1,8 @@
 import logging
+from typing import List, Optional, Dict, Text, Set, Any
+
 import numpy as np
 import scipy.sparse
-from typing import List, Optional, Dict, Text, Set, Any
 
 from rasa.core.featurizers.precomputation import MessageContainerForCoreFeaturization
 from rasa.nlu.extractors.extractor import EntityTagSpec
@@ -360,6 +361,26 @@ class SingleStateFeaturizer:
             for action in domain.action_names_or_texts
         ]
 
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "action_texts": self.action_texts,
+            "entity_tag_specs": self.entity_tag_specs,
+            "feature_states": self._default_feature_states,
+        }
+
+    @classmethod
+    def create_from_dict(
+        cls, data: Dict[str, Any]
+    ) -> Optional["SingleStateFeaturizer"]:
+        if not data:
+            return None
+
+        featurizer = SingleStateFeaturizer()
+        featurizer.action_texts = data["action_texts"]
+        featurizer._default_feature_states = data["feature_states"]
+        featurizer.entity_tag_specs = data["entity_tag_specs"]
+        return featurizer
+
 
 class IntentTokenizerSingleStateFeaturizer(SingleStateFeaturizer):
     """A SingleStateFeaturizer for use with policies that predict intent labels."""

rasa/core/featurizers/tracker_featurizers.py

@@ -1,11 +1,9 @@
 from __future__ import annotations
-from pathlib import Path
-from collections import defaultdict
-from abc import abstractmethod
-import jsonpickle
-import logging
 
-from tqdm import tqdm
+import logging
+from abc import abstractmethod
+from collections import defaultdict
+from pathlib import Path
 from typing import (
     Tuple,
     List,
@@ -18,25 +16,30 @@ from typing import (
     Set,
     DefaultDict,
     cast,
+    Type,
+    Callable,
+    ClassVar,
 )
+
 import numpy as np
+from tqdm import tqdm
 
-from rasa.core.featurizers.single_state_featurizer import SingleStateFeaturizer
-from rasa.core.featurizers.precomputation import MessageContainerForCoreFeaturization
-from rasa.core.exceptions import InvalidTrackerFeaturizerUsageError
 import rasa.shared.core.trackers
 import rasa.shared.utils.io
-from rasa.shared.nlu.constants import TEXT, INTENT, ENTITIES, ACTION_NAME
-from rasa.shared.nlu.training_data.features import Features
-from rasa.shared.core.trackers import DialogueStateTracker
-from rasa.shared.core.domain import State, Domain
-from rasa.shared.core.events import Event, ActionExecuted, UserUttered
+from rasa.core.exceptions import InvalidTrackerFeaturizerUsageError
+from rasa.core.featurizers.precomputation import MessageContainerForCoreFeaturization
+from rasa.core.featurizers.single_state_featurizer import SingleStateFeaturizer
 from rasa.shared.core.constants import (
     USER,
     ACTION_UNLIKELY_INTENT_NAME,
     PREVIOUS_ACTION,
 )
+from rasa.shared.core.domain import State, Domain
+from rasa.shared.core.events import Event, ActionExecuted, UserUttered
+from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.exceptions import RasaException
+from rasa.shared.nlu.constants import TEXT, INTENT, ENTITIES, ACTION_NAME
+from rasa.shared.nlu.training_data.features import Features
 from rasa.utils.tensorflow.constants import LABEL_PAD_ID
 from rasa.utils.tensorflow.model_data import ragged_array_to_ndarray
 
@@ -64,6 +67,10 @@ class InvalidStory(RasaException):
 class TrackerFeaturizer:
     """Base class for actual tracker featurizers."""
 
+    # Class registry to store all subclasses
+    _registry: ClassVar[Dict[str, Type["TrackerFeaturizer"]]] = {}
+    _featurizer_type: str = "TrackerFeaturizer"
+
     def __init__(
         self, state_featurizer: Optional[SingleStateFeaturizer] = None
     ) -> None:
@@ -74,6 +81,36 @@ class TrackerFeaturizer:
         """
         self.state_featurizer = state_featurizer
 
+    @classmethod
+    def register(cls, featurizer_type: str) -> Callable:
+        """Decorator to register featurizer subclasses."""
+
+        def wrapper(subclass: Type["TrackerFeaturizer"]) -> Type["TrackerFeaturizer"]:
+            cls._registry[featurizer_type] = subclass
+            # Store the type identifier in the class for serialization
+            subclass._featurizer_type = featurizer_type
+            return subclass
+
+        return wrapper
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "TrackerFeaturizer":
+        """Create featurizer instance from dictionary."""
+        featurizer_type = data.pop("type")
+
+        if featurizer_type not in cls._registry:
+            raise ValueError(f"Unknown featurizer type: {featurizer_type}")
+
+        # Get the correct subclass and instantiate it
+        subclass = cls._registry[featurizer_type]
+        return subclass.create_from_dict(data)
+
+    @classmethod
+    @abstractmethod
+    def create_from_dict(cls, data: Dict[str, Any]) -> "TrackerFeaturizer":
+        """Each subclass must implement its own creation from dict method."""
+        pass
+
     @staticmethod
     def _create_states(
         tracker: DialogueStateTracker,
@@ -465,9 +502,7 @@ class TrackerFeaturizer:
             self.state_featurizer.entity_tag_specs = []
 
         # noinspection PyTypeChecker
-        rasa.shared.utils.io.write_text_file(
-            str(jsonpickle.encode(self)), featurizer_file
-        )
+        rasa.shared.utils.io.dump_obj_as_json_to_file(featurizer_file, self.to_dict())
 
     @staticmethod
     def load(path: Union[Text, Path]) -> Optional[TrackerFeaturizer]:
@@ -481,7 +516,17 @@ class TrackerFeaturizer:
         """
         featurizer_file = Path(path) / FEATURIZER_FILE
         if featurizer_file.is_file():
-            return jsonpickle.decode(rasa.shared.utils.io.read_file(featurizer_file))
+            data = rasa.shared.utils.io.read_json_file(featurizer_file)
+
+            if "type" not in data:
+                logger.error(
+                    f"Couldn't load featurizer for policy. "
+                    f"File '{featurizer_file}' does not contain all "
+                    f"necessary information. 'type' is missing."
+                )
+                return None
+
+            return TrackerFeaturizer.from_dict(data)
 
         logger.error(
             f"Couldn't load featurizer for policy. "
@@ -508,7 +553,16 @@ class TrackerFeaturizer:
             )
         ]
 
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "type": self.__class__._featurizer_type,
+            "state_featurizer": (
+                self.state_featurizer.to_dict() if self.state_featurizer else None
+            ),
+        }
+
 
+@TrackerFeaturizer.register("FullDialogueTrackerFeaturizer")
 class FullDialogueTrackerFeaturizer(TrackerFeaturizer):
     """Creates full dialogue training data for time distributed architectures.
 
@@ -646,7 +700,20 @@ class FullDialogueTrackerFeaturizer(TrackerFeaturizer):
 
         return trackers_as_states
 
+    def to_dict(self) -> Dict[str, Any]:
+        return super().to_dict()
 
+    @classmethod
+    def create_from_dict(cls, data: Dict[str, Any]) -> "FullDialogueTrackerFeaturizer":
+        state_featurizer = SingleStateFeaturizer.create_from_dict(
+            data["state_featurizer"]
+        )
+        return cls(
+            state_featurizer,
+        )
+
+
+@TrackerFeaturizer.register("MaxHistoryTrackerFeaturizer")
 class MaxHistoryTrackerFeaturizer(TrackerFeaturizer):
     """Truncates the tracker history into `max_history` long sequences.
 
@@ -884,7 +951,25 @@ class MaxHistoryTrackerFeaturizer(TrackerFeaturizer):
 
         return trackers_as_states
 
+    def to_dict(self) -> Dict[str, Any]:
+        data = super().to_dict()
+        data.update(
+            {
+                "remove_duplicates": self.remove_duplicates,
+                "max_history": self.max_history,
+            }
+        )
+        return data
+
+    @classmethod
+    def create_from_dict(cls, data: Dict[str, Any]) -> "MaxHistoryTrackerFeaturizer":
+        state_featurizer = SingleStateFeaturizer.create_from_dict(
+            data["state_featurizer"]
+        )
+        return cls(state_featurizer, data["max_history"], data["remove_duplicates"])
 
+
+@TrackerFeaturizer.register("IntentMaxHistoryTrackerFeaturizer")
 class IntentMaxHistoryTrackerFeaturizer(MaxHistoryTrackerFeaturizer):
     """Truncates the tracker history into `max_history` long sequences.
 
@@ -1159,6 +1244,18 @@ class IntentMaxHistoryTrackerFeaturizer(MaxHistoryTrackerFeaturizer):
 
         return trackers_as_states
 
+    def to_dict(self) -> Dict[str, Any]:
+        return super().to_dict()
+
+    @classmethod
+    def create_from_dict(
+        cls, data: Dict[str, Any]
+    ) -> "IntentMaxHistoryTrackerFeaturizer":
+        state_featurizer = SingleStateFeaturizer.create_from_dict(
+            data["state_featurizer"]
+        )
+        return cls(state_featurizer, data["max_history"], data["remove_duplicates"])
+
 
 def _is_prev_action_unlikely_intent_in_state(state: State) -> bool:
     prev_action_name = state.get(PREVIOUS_ACTION, {}).get(ACTION_NAME)

rasa/core/policies/ted_policy.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
-import logging
 
-from rasa.engine.recipes.default_recipe import DefaultV1Recipe
+import logging
 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,18 +49,22 @@ 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
+from rasa.shared.nlu.training_data.features import (
+    Features,
+    save_features,
+    load_features,
+)
 import rasa.shared.utils.io
 import rasa.utils.io
 from rasa.utils import train_utils
-from rasa.utils.tensorflow.models import RasaModel, TransformerRasaModel
-from rasa.utils.tensorflow import rasa_layers
-from rasa.utils.tensorflow.model_data import (
-    RasaModelData,
-    FeatureSignature,
+from rasa.utils.tensorflow.feature_array import (
     FeatureArray,
-    Data,
+    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_utils import convert_to_data_format
 from rasa.utils.tensorflow.constants import (
     LABEL,
@@ -961,22 +965,32 @@ class TEDPolicy(Policy):
             model_path: Path where model is to be persisted
         """
         model_filename = self._metadata_filename()
-        rasa.utils.io.json_pickle(
-            model_path / f"{model_filename}.priority.pkl", self.priority
-        )
-        rasa.utils.io.pickle_dump(
-            model_path / f"{model_filename}.meta.pkl", self.config
+        rasa.shared.utils.io.dump_obj_as_json_to_file(
+            model_path / f"{model_filename}.priority.json", self.priority
         )
-        rasa.utils.io.pickle_dump(
-            model_path / f"{model_filename}.data_example.pkl", self.data_example
+        rasa.shared.utils.io.dump_obj_as_json_to_file(
+            model_path / f"{model_filename}.meta.json", self.config
         )
-        rasa.utils.io.pickle_dump(
-            model_path / f"{model_filename}.fake_features.pkl", self.fake_features
+        # 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}.label_data.pkl",
+        # 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"),
+        )
+        # save fake features
+        metadata = save_features(
+            self.fake_features, str(model_path / f"{model_filename}.fake_features.st")
+        )
+        rasa.shared.utils.io.dump_obj_as_json_to_file(
+            model_path / f"{model_filename}.fake_features_metadata.json", metadata
         )
+
         entity_tag_specs = (
             [tag_spec._asdict() for tag_spec in self._entity_tag_specs]
             if self._entity_tag_specs
@@ -994,18 +1008,29 @@ class TEDPolicy(Policy):
             model_path: Path where model is to be persisted.
         """
         tf_model_file = model_path / f"{cls._metadata_filename()}.tf_model"
-        loaded_data = rasa.utils.io.pickle_load(
-            model_path / f"{cls._metadata_filename()}.data_example.pkl"
+
+        # 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"),
         )
-        label_data = rasa.utils.io.pickle_load(
-            model_path / f"{cls._metadata_filename()}.label_data.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"),
         )
-        fake_features = rasa.utils.io.pickle_load(
-            model_path / f"{cls._metadata_filename()}.fake_features.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"
         )
-        label_data = RasaModelData(data=label_data)
-        priority = rasa.utils.io.json_unpickle(
-            model_path / f"{cls._metadata_filename()}.priority.pkl"
+        fake_features = load_features(
+            str(model_path / f"{cls._metadata_filename()}.fake_features.st"), metadata
+        )
+
+        priority = rasa.shared.utils.io.read_json_file(
+            model_path / f"{cls._metadata_filename()}.priority.json"
         )
         entity_tag_specs = rasa.shared.utils.io.read_json_file(
             model_path / f"{cls._metadata_filename()}.entity_tag_specs.json"
@@ -1023,8 +1048,8 @@ class TEDPolicy(Policy):
             )
             for tag_spec in entity_tag_specs
         ]
-        model_config = rasa.utils.io.pickle_load(
-            model_path / f"{cls._metadata_filename()}.meta.pkl"
+        model_config = rasa.shared.utils.io.read_json_file(
+            model_path / f"{cls._metadata_filename()}.meta.json"
         )
 
         return {
@@ -1070,7 +1095,7 @@ class TEDPolicy(Policy):
     ) -> TEDPolicy:
         featurizer = TrackerFeaturizer.load(model_path)
 
-        if not (model_path / f"{cls._metadata_filename()}.data_example.pkl").is_file():
+        if not (model_path / f"{cls._metadata_filename()}.data_example.st").is_file():
             return cls(
                 config,
                 model_storage,

rasa/core/policies/unexpected_intent_policy.py

@@ -5,6 +5,7 @@ from typing import Any, List, Optional, Text, Dict, Type, Union
 
 import numpy as np
 import tensorflow as tf
+
 import rasa.utils.common
 from rasa.engine.graph import ExecutionContext
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
@@ -16,6 +17,7 @@ from rasa.shared.core.domain import Domain
 from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.core.constants import SLOTS, ACTIVE_LOOP, ACTION_UNLIKELY_INTENT_NAME
 from rasa.shared.core.events import UserUttered, ActionExecuted
+import rasa.shared.utils.io
 from rasa.shared.nlu.constants import (
     INTENT,
     TEXT,
@@ -103,8 +105,6 @@ from rasa.utils.tensorflow.constants import (
 )
 from rasa.utils.tensorflow import layers
 from rasa.utils.tensorflow.model_data import RasaModelData, FeatureArray, Data
-
-import rasa.utils.io as io_utils
 from rasa.core.exceptions import RasaCoreException
 from rasa.shared.utils import common
 
@@ -881,9 +881,12 @@ class UnexpecTEDIntentPolicy(TEDPolicy):
             model_path: Path where model is to be persisted
         """
         super().persist_model_utilities(model_path)
-        io_utils.pickle_dump(
-            model_path / f"{self._metadata_filename()}.label_quantiles.pkl",
-            self.label_quantiles,
+
+        from safetensors.numpy import save_file
+
+        save_file(
+            {str(k): np.array(v) for k, v in self.label_quantiles.items()},
+            model_path / f"{self._metadata_filename()}.label_quantiles.st",
         )
 
     @classmethod
@@ -894,9 +897,14 @@ class UnexpecTEDIntentPolicy(TEDPolicy):
             model_path: Path where model is to be persisted.
         """
         model_utilties = super()._load_model_utilities(model_path)
-        label_quantiles = io_utils.pickle_load(
-            model_path / f"{cls._metadata_filename()}.label_quantiles.pkl"
+
+        from safetensors.numpy import load_file
+
+        loaded_label_quantiles = load_file(
+            model_path / f"{cls._metadata_filename()}.label_quantiles.st"
         )
+        label_quantiles = {int(k): list(v) for k, v in loaded_label_quantiles.items()}
+
         model_utilties.update({"label_quantiles": label_quantiles})
         return model_utilties
 

rasa/nlu/classifiers/diet_classifier.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
+
 import copy
 import logging
 from collections import defaultdict
 from pathlib import Path
-
-from rasa.exceptions import ModelNotFound
-from rasa.nlu.featurizers.featurizer import Featurizer
+from typing import Any, Dict, List, Optional, Text, Tuple, Union, TypeVar, Type
 
 import numpy as np
 import scipy.sparse
 import tensorflow as tf
 
-from typing import Any, Dict, List, Optional, Text, Tuple, Union, TypeVar, Type
-
+from rasa.exceptions import ModelNotFound
+from rasa.nlu.featurizers.featurizer import Featurizer
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
@@ -20,18 +19,21 @@ 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 (
@@ -118,7 +120,6 @@ LABEL_SUB_KEY = IDS
 
 POSSIBLE_TAGS = [ENTITY_ATTRIBUTE_TYPE, ENTITY_ATTRIBUTE_ROLE, ENTITY_ATTRIBUTE_GROUP]
 
-
 DIETClassifierT = TypeVar("DIETClassifierT", bound="DIETClassifier")
 
 
@@ -1083,18 +1084,24 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
 
             self.model.save(str(tf_model_file))
 
-            io_utils.pickle_dump(
-                model_path / f"{file_name}.data_example.pkl", self._data_example
-            )
-            io_utils.pickle_dump(
-                model_path / f"{file_name}.sparse_feature_sizes.pkl",
-                self._sparse_feature_sizes,
+            # 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}.label_data.pkl",
+            # 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",
             )
-            io_utils.json_pickle(
+
+            rasa.shared.utils.io.dump_obj_as_json_to_file(
+                model_path / f"{file_name}.sparse_feature_sizes.json",
+                self._sparse_feature_sizes,
+            )
+            rasa.shared.utils.io.dump_obj_as_json_to_file(
                 model_path / f"{file_name}.index_label_id_mapping.json",
                 self.index_label_id_mapping,
             )
@@ -1183,15 +1190,22 @@ class DIETClassifier(GraphComponent, IntentClassifier, EntityExtractorMixin):
     ]:
         file_name = cls.__name__
 
-        data_example = io_utils.pickle_load(
-            model_path / f"{file_name}.data_example.pkl"
+        # 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"),
         )
-        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"
+        # 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"
         )
-        index_label_id_mapping = io_utils.json_unpickle(
+        index_label_id_mapping = rasa.shared.utils.io.read_json_file(
             model_path / f"{file_name}.index_label_id_mapping.json"
         )
         entity_tag_specs = rasa.shared.utils.io.read_json_file(
@@ -1211,7 +1225,6 @@ 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()
         }