rasa-pro 3.12.0.dev2__py3-none-any.whl → 3.12.0.dev3__py3-none-any.whl

rasa/tracing/config.py

@@ -30,8 +30,8 @@ from rasa.dialogue_understanding.commands import (
 )
 from rasa.dialogue_understanding.generator import (
     LLMCommandGenerator,
-    SingleStepLLMCommandGenerator,
     MultiStepLLMCommandGenerator,
+    SingleStepLLMCommandGenerator,
 )
 from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.dialogue_understanding.generator.nlu_command_adapter import NLUCommandAdapter
@@ -40,11 +40,11 @@ from rasa.engine.training.graph_trainer import GraphTrainer
 from rasa.tracing.constants import (
     ENDPOINTS_ENDPOINT_KEY,
     ENDPOINTS_INSECURE_KEY,
+    ENDPOINTS_METRICS_KEY,
     ENDPOINTS_OTLP_BACKEND_TYPE,
+    ENDPOINTS_ROOT_CERTIFICATES_KEY,
     ENDPOINTS_TRACING_KEY,
     ENDPOINTS_TRACING_SERVICE_NAME_KEY,
-    ENDPOINTS_ROOT_CERTIFICATES_KEY,
-    ENDPOINTS_METRICS_KEY,
 )
 from rasa.tracing.instrumentation import instrumentation
 from rasa.tracing.metric_instrument_provider import MetricInstrumentProvider

rasa/tracing/instrumentation/attribute_extractors.py

@@ -578,6 +578,7 @@ def extract_attrs_for_run_step(
     tracker: DialogueStateTracker,
     available_actions: List[str],
     flows: FlowsList,
+    previous_step_id: Text,
 ) -> Dict[str, Any]:
     current_context = extract_current_context_attribute(stack)
 
@@ -586,6 +587,7 @@ def extract_attrs_for_run_step(
         "step_description": step.description if step.description else "None",
         "current_flow_id": flow.id,
         "current_context": json.dumps(current_context),
+        "previous_step_id": previous_step_id,
     }
 
 

rasa/tracing/instrumentation/instrumentation.py

@@ -24,7 +24,7 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.trace import SpanKind, Tracer
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
-from rasa.core.actions.action import Action, RemoteAction, CustomActionExecutor
+from rasa.core.actions.action import Action, CustomActionExecutor, RemoteAction
 from rasa.core.actions.custom_action_executor import RetryCustomActionExecutor
 from rasa.core.actions.grpc_custom_action_executor import GRPCCustomActionExecutor
 from rasa.core.agent import Agent
@@ -63,13 +63,13 @@ from rasa.tracing.instrumentation.intentless_policy_instrumentation import (
     _instrument_select_response_examples,
 )
 from rasa.tracing.instrumentation.metrics import (
+    record_callable_duration_metrics,
     record_llm_command_generator_metrics,
-    record_single_step_llm_command_generator_metrics,
     record_multi_step_llm_command_generator_metrics,
-    record_callable_duration_metrics,
     record_request_size_in_bytes,
+    record_single_step_llm_command_generator_metrics,
 )
-from rasa.utils.endpoints import concat_url, EndpointConfig
+from rasa.utils.endpoints import EndpointConfig, concat_url
 
 # The `TypeVar` representing the return type for a function to be wrapped.
 S = TypeVar("S")

rasa/tracing/instrumentation/intentless_policy_instrumentation.py

@@ -1,7 +1,7 @@
 import functools
 import json
 import logging
-from typing import Callable, List, Optional, TYPE_CHECKING, Type
+from typing import TYPE_CHECKING, Callable, List, Optional, Type
 
 from opentelemetry.trace import Tracer
 

rasa/tracing/instrumentation/metrics.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Any, Dict
 
 import psutil
 
@@ -7,29 +7,29 @@ from rasa.core.policies.enterprise_search_policy import EnterpriseSearchPolicy
 from rasa.core.policies.intentless_policy import IntentlessPolicy
 from rasa.dialogue_understanding.generator import (
     LLMCommandGenerator,
-    SingleStepLLMCommandGenerator,
     MultiStepLLMCommandGenerator,
+    SingleStepLLMCommandGenerator,
 )
 from rasa.tracing.constants import (
+    CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
+    ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
+    INTENTLESS_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
     LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     PROMPT_TOKEN_LENGTH_ATTRIBUTE_NAME,
-    LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
-    ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
-    INTENTLESS_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
-    CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
     RASA_CLIENT_REQUEST_BODY_SIZE_METRIC_NAME,
     RASA_CLIENT_REQUEST_DURATION_METRIC_NAME,
     REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
 )
 from rasa.tracing.metric_instrument_provider import MetricInstrumentProvider
 from rasa.utils.endpoints import EndpointConfig

rasa/tracing/metric_instrument_provider.py

@@ -1,30 +1,30 @@
-from typing import ClassVar, Dict, Any
+from typing import Any, ClassVar, Dict
 
 from opentelemetry.metrics import get_meter_provider
 from opentelemetry.sdk.metrics import Meter
 
-from rasa.utils.singleton import Singleton
 from rasa.tracing.constants import (
+    CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
+    DURATION_UNIT_NAME,
+    ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
+    INTENTLESS_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
+    LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
     LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
-    LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
-    DURATION_UNIT_NAME,
-    ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
-    INTENTLESS_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
-    CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
-    RASA_CLIENT_REQUEST_DURATION_METRIC_NAME,
     RASA_CLIENT_REQUEST_BODY_SIZE_METRIC_NAME,
+    RASA_CLIENT_REQUEST_DURATION_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
 )
+from rasa.utils.singleton import Singleton
 
 
 class MetricInstrumentProvider(metaclass=Singleton):

rasa/utils/common.py

@@ -8,38 +8,39 @@ import shutil
 import tempfile
 import warnings
 from pathlib import Path
+from socket import SOCK_DGRAM, SOCK_STREAM
 from types import TracebackType
 from typing import (
     Any,
+    ContextManager,
     Coroutine,
     Dict,
     List,
     Optional,
+    Set,
     Text,
+    Tuple,
     Type,
     TypeVar,
     Union,
-    ContextManager,
-    Set,
-    Tuple,
 )
 
-from socket import SOCK_DGRAM, SOCK_STREAM
 import numpy as np
+
+import rasa.shared.utils.io
 import rasa.utils.io
 from rasa.anonymization import ENV_LOG_LEVEL_FAKER, ENV_LOG_LEVEL_PRESIDIO
 from rasa.constants import (
     DEFAULT_LOG_LEVEL_LIBRARIES,
+    ENV_LOG_LEVEL_KAFKA,
     ENV_LOG_LEVEL_LIBRARIES,
     ENV_LOG_LEVEL_MATPLOTLIB,
     ENV_LOG_LEVEL_MLFLOW,
     ENV_LOG_LEVEL_RABBITMQ,
-    ENV_LOG_LEVEL_KAFKA,
 )
 from rasa.shared.constants import DEFAULT_LOG_LEVEL, ENV_LOG_LEVEL, TCP_PROTOCOL
 from rasa.shared.exceptions import RasaException
-import rasa.shared.utils.io
-from rasa.shared.utils.yaml import read_yaml_file, write_yaml, read_config_file
+from rasa.shared.utils.yaml import read_config_file, read_yaml_file, write_yaml
 
 logger = logging.getLogger(__name__)
 
@@ -146,7 +147,9 @@ def read_global_config(path: Text) -> Dict[Text, Any]:
     """
     # noinspection PyBroadException
     try:
-        return read_config_file(path)
+        # set expand_env_vars to False because this is a global config file
+        # used for telemetry, in which case we don't need to expand env vars
+        return read_config_file(path, expand_env_vars=False)
     except Exception:
         # if things go south we pretend there is no config
         return {}
@@ -299,7 +302,7 @@ def update_sanic_log_level(
     syslog_protocol: Optional[Text] = None,
 ) -> None:
     """Set the log level to 'LOG_LEVEL_LIBRARIES' environment variable ."""
-    from sanic.log import logger, error_logger, access_logger
+    from sanic.log import access_logger, error_logger, logger
 
     log_level = os.environ.get(ENV_LOG_LEVEL_LIBRARIES, DEFAULT_LOG_LEVEL_LIBRARIES)
 

rasa/utils/endpoints.py

@@ -1,7 +1,7 @@
 import os
 import ssl
 from types import ModuleType
-from typing import Any, List, Optional, Text, Dict, Union
+from typing import Any, Dict, List, Optional, Text, Union
 
 import aiohttp
 import structlog

rasa/utils/io.py

@@ -2,22 +2,22 @@ import asyncio
 import filecmp
 import logging
 import os
+import re
 import tempfile
 import warnings
-import re
 from asyncio import AbstractEventLoop
 from collections import OrderedDict
 from io import StringIO
 from pathlib import Path
 from typing import (
-    Optional,
-    Text,
+    TYPE_CHECKING,
     Any,
-    Union,
+    Callable,
     List,
+    Optional,
+    Text,
     Type,
-    Callable,
-    TYPE_CHECKING,
+    Union,
 )
 
 from ruamel import yaml
@@ -153,8 +153,8 @@ def create_validator(
     """Helper method to create `Validator` classes from callable functions. Should be
     removed when questionary supports `Validator` objects.
     """
-    from prompt_toolkit.validation import Validator, ValidationError
     from prompt_toolkit.document import Document
+    from prompt_toolkit.validation import ValidationError, Validator
 
     class FunctionValidator(Validator):
         @staticmethod

rasa/utils/licensing.py

@@ -1,4 +1,3 @@
-from asyncio import AbstractEventLoop
 import hashlib
 import os
 import random
@@ -6,19 +5,19 @@ import re
 import time
 import typing
 import uuid
+from asyncio import AbstractEventLoop
 from datetime import datetime, timezone
 from typing import Any, Callable, Dict, Optional, Set, Text, TypeVar
 
 import jwt
+import structlog
 from dotenv import dotenv_values
 from sanic import Sanic
-import structlog
-from rasa import telemetry
 
+from rasa import telemetry
 from rasa.core import jobs
 from rasa.shared.utils.cli import print_error_and_exit
 
-
 if typing.TYPE_CHECKING:
     from rasa.core.tracker_store import TrackerStore
 

rasa/utils/log_utils.py

@@ -1,22 +1,23 @@
 from __future__ import annotations
-import os
+
 import logging
+import os
 import sys
 from typing import Any, Dict, Optional
 
 import structlog
-from structlog_sentry import SentryProcessor
 from structlog.dev import ConsoleRenderer
 from structlog.typing import EventDict, WrappedLogger
+from structlog_sentry import SentryProcessor
+
+from rasa.plugin import plugin_manager
 from rasa.shared.constants import (
-    ENV_LOG_LEVEL,
     DEFAULT_LOG_LEVEL,
+    DEFAULT_LOG_LEVEL_LLM,
+    ENV_LOG_LEVEL,
     ENV_LOG_LEVEL_LLM,
     ENV_LOG_LEVEL_LLM_MODULE_NAMES,
-    DEFAULT_LOG_LEVEL_LLM,
 )
-from rasa.plugin import plugin_manager
-
 
 FORCE_JSON_LOGGING = os.environ.get("FORCE_JSON_LOGGING")
 

rasa/utils/ml_utils.py

@@ -5,6 +5,7 @@ import numpy as np
 import structlog
 from langchain.schema.embeddings import Embeddings
 from langchain_community.vectorstores import FAISS
+
 from rasa.shared.constants import REQUIRED_SLOTS_KEY
 from rasa.shared.core.domain import KEY_RESPONSES_TEXT, Domain
 from rasa.shared.utils.llm import AI

rasa/utils/plotting.py

@@ -1,11 +1,11 @@
-import logging
 import itertools
+import logging
 import os
 from functools import wraps
+from typing import Any, Callable, List, Optional, Text, Tuple, TypeVar, Union
 
-import numpy as np
-from typing import Any, Callable, List, Optional, Text, TypeVar, Union, Tuple
 import matplotlib
+import numpy as np
 from matplotlib.ticker import FormatStrFormatter
 
 import rasa.shared.utils.io

rasa/utils/sanic_error_handler.py

@@ -1,7 +1,7 @@
 from sanic import Sanic
+from sanic.exceptions import ServerError
 from sanic.handlers import ErrorHandler
 from sanic.request import Request
-from sanic.exceptions import ServerError
 
 
 # TODO: remove custom handler when upgrading to sanic >= 24

rasa/utils/tensorflow/callback.py

@@ -1,7 +1,7 @@
+import logging
 from pathlib import Path
-from typing import Dict, Text, Any, Optional
+from typing import Any, Dict, Optional, Text
 
-import logging
 import tensorflow as tf
 from tqdm import tqdm
 

rasa/utils/tensorflow/crf.py

@@ -1,8 +1,8 @@
+from typing import Any, List, Optional, Tuple, Union
+
 import tensorflow as tf
 from tensorflow import TensorShape
 from tensorflow.types.experimental import TensorLike
-from typing import Tuple, Any, List, Union, Optional
-
 
 # original code taken from
 # https://github.com/tensorflow/addons/blob/b8cab7fd61af4f697a1cdae4f51c37c346b9c6f0/tensorflow_addons/text/crf.py

rasa/utils/tensorflow/data_generator.py

@@ -1,13 +1,13 @@
+import logging
 import math
-from typing import List, Union, Text, Optional, Any, Tuple, Dict, cast
+from typing import Any, Dict, List, Optional, Text, Tuple, Union, cast
 
-import logging
-import scipy.sparse
 import numpy as np
+import scipy.sparse
 from tensorflow.keras.utils import Sequence
 
-from rasa.utils.tensorflow.constants import SEQUENCE, BALANCED
-from rasa.utils.tensorflow.model_data import RasaModelData, Data, FeatureArray
+from rasa.utils.tensorflow.constants import BALANCED, SEQUENCE
+from rasa.utils.tensorflow.model_data import Data, FeatureArray, RasaModelData
 
 logger = logging.getLogger(__name__)
 

rasa/utils/tensorflow/environment.py

@@ -1,16 +1,16 @@
 import logging
 import os
-from typing import Text, Dict
 import typing
+from typing import Dict, Text
 
 import rasa.shared.utils.io
 from rasa.constants import (
-    ENV_GPU_CONFIG,
     ENV_CPU_INTER_OP_CONFIG,
     ENV_CPU_INTRA_OP_CONFIG,
+    ENV_GPU_CONFIG,
 )
-from rasa.utils.tensorflow.constants import TF_DETERMINISTIC_OPS
 from rasa.shared.utils import io as shared_io_utils
+from rasa.utils.tensorflow.constants import TF_DETERMINISTIC_OPS
 
 if typing.TYPE_CHECKING:
     from tensorflow import config as tf_config

rasa/utils/tensorflow/feature_array.py

@@ -1,9 +1,8 @@
-from typing import Dict, Any, List, Tuple, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
-from safetensors.numpy import save_file
 import numpy as np
-from safetensors.numpy import load_file
 import scipy.sparse
+from safetensors.numpy import load_file, save_file
 
 import rasa.shared.utils.io
 

rasa/utils/tensorflow/layers.py

@@ -1,29 +1,35 @@
 import logging
-from typing import List, Optional, Text, Tuple, Callable, Union, Any
+from typing import Any, Callable, List, Optional, Text, Tuple, Union
+
 import tensorflow as tf
+import tensorflow.keras.backend as K
 
 # TODO: The following is not (yet) available via tf.keras
 from keras.src.utils.control_flow_util import smart_cond
-import tensorflow.keras.backend as K
 
 import rasa.utils.tensorflow.crf
+import rasa.utils.tensorflow.layers_utils as layers_utils
+from rasa.core.constants import DIALOGUE
+from rasa.shared.nlu.constants import (
+    ACTION_NAME,
+    ACTION_TEXT,
+    FEATURE_TYPE_SENTENCE,
+    FEATURE_TYPE_SEQUENCE,
+    INTENT,
+    TEXT,
+)
 from rasa.utils.tensorflow.constants import (
-    SOFTMAX,
-    MARGIN,
     COSINE,
-    INNER,
     CROSS_ENTROPY,
+    INNER,
     LABEL,
     LABEL_PAD_ID,
+    MARGIN,
+    SOFTMAX,
 )
-from rasa.core.constants import DIALOGUE
-from rasa.shared.nlu.constants import FEATURE_TYPE_SENTENCE, FEATURE_TYPE_SEQUENCE
-from rasa.shared.nlu.constants import TEXT, INTENT, ACTION_NAME, ACTION_TEXT
-
-from rasa.utils.tensorflow.metrics import F1Score
-from rasa.utils.tensorflow.exceptions import TFLayerConfigException
-import rasa.utils.tensorflow.layers_utils as layers_utils
 from rasa.utils.tensorflow.crf import crf_log_likelihood
+from rasa.utils.tensorflow.exceptions import TFLayerConfigException
+from rasa.utils.tensorflow.metrics import F1Score
 
 logger = logging.getLogger(__name__)
 

rasa/utils/tensorflow/layers_utils.py

@@ -1,6 +1,7 @@
+from typing import Optional, Union
+
 import tensorflow as tf
 from tensorflow import Tensor
-from typing import Union, Optional
 
 
 def random_indices(

rasa/utils/tensorflow/metrics.py

@@ -1,8 +1,8 @@
+from typing import Any, Dict, Optional
+
 import tensorflow as tf
 from tensorflow.keras import backend as K
 from tensorflow.types.experimental import TensorLike
-from typing import Any, Dict, Optional
-
 
 # original code taken from
 # https://github.com/tensorflow/addons/blob/f30df4322b5580b3e5946530a60f7126035dd73b/tensorflow_addons/metrics/f_scores.py

rasa/utils/tensorflow/model_data.py

@@ -1,20 +1,20 @@
 import logging
+from collections import OrderedDict, defaultdict
 from typing import (
-    Optional,
+    Any,
     DefaultDict,
     Dict,
+    ItemsView,
     Iterable,
-    Text,
     List,
+    NamedTuple,
+    Optional,
+    Text,
     Tuple,
-    Any,
     Union,
-    NamedTuple,
-    ItemsView,
-    overload,
     cast,
+    overload,
 )
-from collections import defaultdict, OrderedDict
 
 import numpy as np
 import scipy.sparse

rasa/utils/tensorflow/model_data_utils.py

@@ -1,25 +1,26 @@
-import typing
 import copy
+import typing
+from collections import OrderedDict, defaultdict
+from typing import Any, DefaultDict, Dict, List, Optional, Text, Tuple, Union, cast
+
 import numpy as np
 import scipy.sparse
-from collections import defaultdict, OrderedDict
-from typing import List, Optional, Text, Dict, Tuple, Union, Any, DefaultDict, cast
 
 from rasa.nlu.constants import TOKENS_NAMES
-from rasa.utils.tensorflow.model_data import Data, FeatureArray, ragged_array_to_ndarray
-from rasa.utils.tensorflow.constants import MASK, IDS
-from rasa.shared.nlu.training_data.message import Message
 from rasa.shared.nlu.constants import (
-    TEXT,
     ENTITIES,
-    ENTITY_ATTRIBUTE_TYPE,
     ENTITY_ATTRIBUTE_GROUP,
     ENTITY_ATTRIBUTE_ROLE,
+    ENTITY_ATTRIBUTE_TYPE,
+    TEXT,
 )
+from rasa.shared.nlu.training_data.message import Message
+from rasa.utils.tensorflow.constants import IDS, MASK
+from rasa.utils.tensorflow.model_data import Data, FeatureArray, ragged_array_to_ndarray
 
 if typing.TYPE_CHECKING:
-    from rasa.shared.nlu.training_data.features import Features
     from rasa.nlu.extractors.extractor import EntityTagSpec
+    from rasa.shared.nlu.training_data.features import Features
 
 TAG_ID_ORIGIN = "tag_id_origin"
 

rasa/utils/tensorflow/models.py

@@ -1,53 +1,52 @@
-import time
-import random
-import tensorflow as tf
-import numpy as np
 import logging
 import os
+import random
+import time
 from collections import defaultdict
-from typing import List, Text, Dict, Tuple, Union, Optional, Any, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text, Tuple, Union
 
-from keras.src.utils import tf_utils
+import numpy as np
+import tensorflow as tf
 from keras import Model
+from keras.src.utils import tf_utils
 
+import rasa.utils.train_utils
 from rasa.shared.constants import DIAGNOSTIC_DATA
+from rasa.shared.exceptions import RasaException
+from rasa.shared.nlu.constants import TEXT
+from rasa.utils.tensorflow import layers, rasa_layers
 from rasa.utils.tensorflow.constants import (
-    LABEL,
+    CONNECTION_DENSITY,
+    CONSTRAIN_SIMILARITIES,
+    EMBEDDING_DIMENSION,
     IDS,
     INTENT_CLASSIFICATION,
-    SENTENCE,
-    SEQUENCE_LENGTH,
-    RANDOM_SEED,
-    EMBEDDING_DIMENSION,
-    REGULARIZATION_CONSTANT,
-    SIMILARITY_TYPE,
-    CONNECTION_DENSITY,
-    NUM_NEG,
+    LABEL,
+    LEARNING_RATE,
     LOSS_TYPE,
-    MAX_POS_SIM,
     MAX_NEG_SIM,
-    USE_MAX_NEG_SIM,
-    NEGATIVE_MARGIN_SCALE,
-    SCALE_LOSS,
-    LEARNING_RATE,
-    CONSTRAIN_SIMILARITIES,
+    MAX_POS_SIM,
     MODEL_CONFIDENCE,
+    NEGATIVE_MARGIN_SCALE,
+    NUM_NEG,
+    RANDOM_SEED,
+    REGULARIZATION_CONSTANT,
     RUN_EAGERLY,
+    SCALE_LOSS,
+    SENTENCE,
+    SEQUENCE_LENGTH,
+    SIMILARITY_TYPE,
+    USE_MAX_NEG_SIM,
 )
-from rasa.utils.tensorflow.model_data import (
-    RasaModelData,
-    FeatureSignature,
-    FeatureArray,
-)
-import rasa.utils.train_utils
-from rasa.utils.tensorflow import layers
-from rasa.utils.tensorflow import rasa_layers
 from rasa.utils.tensorflow.data_generator import (
-    RasaDataGenerator,
     RasaBatchDataGenerator,
+    RasaDataGenerator,
+)
+from rasa.utils.tensorflow.model_data import (
+    FeatureArray,
+    FeatureSignature,
+    RasaModelData,
 )
-from rasa.shared.nlu.constants import TEXT
-from rasa.shared.exceptions import RasaException
 from rasa.utils.tensorflow.types import BatchData, MaybeNestedBatchData
 
 if TYPE_CHECKING: