rasa-pro 3.10.8__tar.gz → 3.10.9.dev1__tar.gz

{rasa_pro-3.10.8 → rasa_pro-3.10.9.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: rasa-pro
-Version: 3.10.8
+Version: 3.10.9.dev1
 Summary: State-of-the-art open-core Conversational AI framework for Enterprises that natively leverages generative AI for effortless assistant development.
 Home-page: https://rasa.com
 Keywords: nlp,machine-learning,machine-learning-library,bot,bots,botkit,rasa conversational-agents,conversational-ai,chatbot,chatbot-framework,bot-framework
@@ -34,7 +34,6 @@ Requires-Dist: attrs (>=23.1,<23.2)
 Requires-Dist: azure-storage-blob (>=12.16.0,<12.17.0)
 Requires-Dist: boto3 (>=1.35.5,<1.36.0)
 Requires-Dist: certifi (>=2024.07.04)
-Requires-Dist: cloudpickle (>=2.2.1,<3.1)
 Requires-Dist: colorama (>=0.4.6,<0.5.0) ; sys_platform == "win32"
 Requires-Dist: colorclass (>=2.2,<2.3)
 Requires-Dist: coloredlogs (>=15,<16)
@@ -57,7 +56,6 @@ Requires-Dist: importlib-metadata (>=8.5.0,<8.6.0)
 Requires-Dist: importlib-resources (==6.1.3)
 Requires-Dist: jieba (>=0.42.1,<0.43) ; extra == "jieba" or extra == "full"
 Requires-Dist: jinja2 (>=3.1.4,<4.0.0)
-Requires-Dist: joblib (>=1.2.0,<1.3.0)
 Requires-Dist: jsonpatch (>=1.33,<2.0)
 Requires-Dist: jsonpickle (>=3.0,<3.1)
 Requires-Dist: jsonschema (>=4.22)
@@ -110,6 +108,7 @@ Requires-Dist: requests (>=2.31.0,<2.32.0)
 Requires-Dist: rich (>=13.4.2,<14.0.0)
 Requires-Dist: rocketchat_API (>=1.30.0,<1.31.0)
 Requires-Dist: ruamel.yaml (>=0.17.21,<0.17.22)
+Requires-Dist: safetensors (>=0.4.5,<0.5.0)
 Requires-Dist: sanic (>=22.12,<22.13)
 Requires-Dist: sanic-cors (>=2.2.0,<2.3.0)
 Requires-Dist: sanic-jwt (>=1.8.0,<2.0.0)
@@ -120,6 +119,7 @@ Requires-Dist: sentencepiece[sentencepiece] (>=0.1.99,<0.2.0) ; extra == "transf
 Requires-Dist: sentry-sdk (>=1.14.0,<1.15.0)
 Requires-Dist: setuptools (>=70.0.0,<70.1.0)
 Requires-Dist: sklearn-crfsuite (>=0.3.6,<0.4.0)
+Requires-Dist: skops (>=0.10.0,<0.11.0)
 Requires-Dist: slack-sdk (>=3.27.1,<4.0.0)
 Requires-Dist: spacy (>=3.5.4,<4.0.0) ; extra == "spacy" or extra == "full"
 Requires-Dist: structlog (>=23.1.0,<23.2.0)

{rasa_pro-3.10.8 → rasa_pro-3.10.9.dev1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rasa-pro"
-version = "3.10.8"
+version = "3.10.9.dev1"
 description = "State-of-the-art open-core Conversational AI framework for Enterprises that natively leverages generative AI for effortless assistant development."
 authors = [ "Rasa Technologies GmbH <hi@rasa.com>",]
 maintainers = [ "Tom Bocklisch <tom@rasa.com>",]
@@ -120,7 +120,6 @@ sanic-cors = "~2.2.0"
 sanic-jwt = "^1.8.0"
 sanic-routing = "^22.8.0"
 websockets = ">=10.4,<11.0"
-cloudpickle = ">=2.2.1,<3.1"
 aiohttp = ">=3.9.4,<3.10"
 questionary = ">=1.10.0,<2.1.0"
 prompt-toolkit = "^3.0.28,<3.0.29"
@@ -136,7 +135,6 @@ tensorflow_hub = "^0.13.0"
 setuptools = "~70.0.0"
 ujson = ">=5.8,<6.0"
 regex = ">=2022.10.31,<2022.11"
-joblib = ">=1.2.0,<1.3.0"
 sentry-sdk = ">=1.14.0,<1.15.0"
 aio-pika = ">=8.2.3,<8.2.4"
 aiogram = ">=2.15,<2.26"
@@ -162,7 +160,7 @@ wheel = ">=0.40.0"
 certifi = ">=2024.07.04"
 cryptography = ">=42.0.5"
 jsonpatch = "^1.33"
-importlib-resources = "6.1.3" # pinned to avoid breaking changes in versions >= 6.2.0 with Python 3.9
+importlib-resources = "6.1.3"
 importlib-metadata = "~8.5.0"
 opentelemetry-sdk = "~1.16.0"
 opentelemetry-exporter-jaeger = "~1.16.0"
@@ -180,16 +178,18 @@ python-keycloak = "^3.12.0"
 tiktoken = "~0.7.0"
 faiss-cpu = "^1.7.4"
 pycountry = "^22.3.5"
-pymilvus = "<2.4.2" # pinned because 2.4.2 introduces milvus-lite dependency that fails on windows
+pymilvus = "<2.4.2"
 qdrant-client = "^1.9.0"
 psutil = "^5.9.5"
-tenacity = "~8.4.1" # pinned down because of langchain
+tenacity = "~8.4.1"
 types-protobuf = "4.25.0.20240417"
 litellm = "~1.50.0"
 openpyxl = "^3.1.5"
 diskcache = "~5.6.3"
+skops = "~0.10.0"
+safetensors = "~0.4.5"
 [[tool.poetry.dependencies.tensorflow-io-gcs-filesystem]]
-version = "==0.31" # newer versions doesn't provide wheels for Windows
+version = "==0.31"
 markers = "sys_platform == 'win32'"
 
 [[tool.poetry.dependencies.tensorflow-io-gcs-filesystem]]
@@ -245,20 +245,6 @@ mlflow = [ "mlflow",]
 rasa = "rasa.__main__:main"
 
 [tool.ruff.lint]
-# The following checks are disabled because they are not relevant for Rasa
-# - RUF001: Checks for ambiguous Unicode characters in strings. https://docs.astral.sh/ruff/rules/ambiguous-unicode-character-string/
-# - RUF002: Checks for ambiguous Unicode characters in docstrings. https://docs.astral.sh/ruff/rules/ambiguous-unicode-character-docstring/
-# - RUF003: Checks for ambiguous Unicode characters in comments. https://docs.astral.sh/ruff/rules/ambiguous-unicode-character-comment/
-# - RUF005: Use of + operator for string concatenation https://docs.astral.sh/ruff/rules/collection-literal-concatenation/
-# Following checks are currently disabled but will be enabled in the future
-# as part of the ticket https://rasahq.atlassian.net/browse/ATO-866
-# - D100: Missing docstring in public module
-# - D101: Missing docstring in public class
-# - D102: Missing docstring in public method
-# - D103: Missing docstring in public function
-# - D104: Missing docstring in public package
-# - D105: Missing docstring in magic method
-# - D107: Missing docstring in __init__
 ignore = [ "D100", "D101", "D102", "D103", "D104", "D105", "D107", "RUF001", "RUF002", "RUF003", "RUF005",]
 select = [ "D", "E", "F", "W", "RUF",]
 

{rasa_pro-3.10.8 → rasa_pro-3.10.9.dev1}/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.10.0rc1"
+MINIMUM_COMPATIBLE_VERSION = "3.10.9.dev1"
 
 GLOBAL_USER_CONFIG_PATH = os.path.expanduser("~/.config/rasa/global.yml")
 

{rasa_pro-3.10.8 → rasa_pro-3.10.9.dev1}/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_pro-3.10.8 → rasa_pro-3.10.9.dev1}/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_pro-3.10.8 → rasa_pro-3.10.9.dev1}/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_pro-3.10.8 → rasa_pro-3.10.9.dev1}/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