rasa-pro 3.8.16__py3-none-any.whl

rasa/utils/tensorflow/model_data.py

@@ -0,0 +1,991 @@
+import logging
+from typing import (
+    Optional,
+    DefaultDict,
+    Dict,
+    Iterable,
+    Text,
+    List,
+    Tuple,
+    Any,
+    Union,
+    NamedTuple,
+    ItemsView,
+    overload,
+    cast,
+)
+from collections import defaultdict, OrderedDict
+
+import numpy as np
+import scipy.sparse
+from sklearn.model_selection import train_test_split
+
+logger = logging.getLogger(__name__)
+
+
+def ragged_array_to_ndarray(ragged_array: Iterable[np.ndarray]) -> np.ndarray:
+    """Converts ragged array to numpy array.
+
+    Ragged array, also known as a jagged array, irregular array is an array of
+    arrays of which the member arrays can be of different lengths.
+    Try to convert as is (preserves type), if it fails because not all numpy arrays have
+    the same shape, then creates numpy array of objects.
+    """
+    try:
+        return np.array(ragged_array)
+    except ValueError:
+        return np.array(ragged_array, dtype=object)
+
+
+class FeatureArray(np.ndarray):
+    """Stores any kind of features ready to be used by a RasaModel.
+
+    Next to the input numpy array of features, it also received the number of
+    dimensions of the features.
+    As our features can have 1 to 4 dimensions we might have different number of numpy
+    arrays stacked. The number of dimensions helps us to figure out how to handle this
+    particular feature array. Also, it is automatically determined whether the feature
+    array is sparse or not and the number of units is determined as well.
+
+    Subclassing np.array: https://numpy.org/doc/stable/user/basics.subclassing.html
+    """
+
+    def __new__(
+        cls, input_array: np.ndarray, number_of_dimensions: int
+    ) -> "FeatureArray":
+        """Create and return a new object.  See help(type) for accurate signature."""
+        FeatureArray._validate_number_of_dimensions(number_of_dimensions, input_array)
+
+        feature_array = np.asarray(input_array).view(cls)
+
+        if number_of_dimensions <= 2:
+            feature_array.units = input_array.shape[-1]
+            feature_array.is_sparse = isinstance(input_array[0], scipy.sparse.spmatrix)
+        elif number_of_dimensions == 3:
+            feature_array.units = input_array[0].shape[-1]
+            feature_array.is_sparse = isinstance(input_array[0], scipy.sparse.spmatrix)
+        elif number_of_dimensions == 4:
+            feature_array.units = input_array[0][0].shape[-1]
+            feature_array.is_sparse = isinstance(
+                input_array[0][0], scipy.sparse.spmatrix
+            )
+        else:
+            raise ValueError(
+                f"Number of dimensions '{number_of_dimensions}' currently not "
+                f"supported."
+            )
+
+        feature_array.number_of_dimensions = number_of_dimensions
+
+        return feature_array
+
+    def __init__(
+        self, input_array: Any, number_of_dimensions: int, **kwargs: Any
+    ) -> None:
+        """Initialize. FeatureArray.
+
+        Needed in order to avoid 'Invalid keyword argument number_of_dimensions
+        to function FeatureArray.__init__ '
+        Args:
+            input_array: the array that contains features
+            number_of_dimensions: number of dimensions in input_array
+        """
+        super().__init__(**kwargs)
+        self.number_of_dimensions = number_of_dimensions
+
+    def __array_finalize__(self, obj: Optional[np.ndarray]) -> None:
+        """This method is called when the system allocates a new array from obj.
+
+        Args:
+            obj: A subclass (subtype) of ndarray.
+        """
+        if obj is None:
+            return
+
+        self.units = getattr(obj, "units", None)
+        self.number_of_dimensions = getattr(obj, "number_of_dimensions", None)  # type: ignore[assignment] # noqa:E501
+        self.is_sparse = getattr(obj, "is_sparse", None)
+
+        default_attributes = {
+            "units": self.units,
+            "number_of_dimensions": self.number_of_dimensions,
+            "is_spare": self.is_sparse,
+        }
+        self.__dict__.update(default_attributes)
+
+    # pytype: disable=attribute-error
+    def __array_ufunc__(
+        self, ufunc: Any, method: Text, *inputs: Any, **kwargs: Any
+    ) -> Any:
+        """Overwrite this method as we are subclassing numpy array.
+
+        Args:
+            ufunc: The ufunc object that was called.
+            method: A string indicating which Ufunc method was called
+                    (one of "__call__", "reduce", "reduceat", "accumulate", "outer",
+                    "inner").
+            *inputs: A tuple of the input arguments to the ufunc.
+            **kwargs: Any additional arguments
+
+        Returns:
+            The result of the operation.
+        """
+        f = {
+            "reduce": ufunc.reduce,
+            "accumulate": ufunc.accumulate,
+            "reduceat": ufunc.reduceat,
+            "outer": ufunc.outer,
+            "at": ufunc.at,
+            "__call__": ufunc,
+        }
+        # convert the inputs to np.ndarray to prevent recursion, call the function,
+        # then cast it back as FeatureArray
+        output = FeatureArray(
+            f[method](*(i.view(np.ndarray) for i in inputs), **kwargs),
+            number_of_dimensions=kwargs["number_of_dimensions"],
+        )
+        output.__dict__ = self.__dict__  # carry forward attributes
+        return output
+
+    def __reduce__(self) -> Tuple[Any, Any, Any]:
+        """Needed in order to pickle this object.
+
+        Returns:
+            A tuple.
+        """
+        pickled_state = super(FeatureArray, self).__reduce__()
+        if isinstance(pickled_state, str):
+            raise TypeError("np array __reduce__ returned string instead of tuple.")
+        new_state = pickled_state[2] + (
+            self.number_of_dimensions,
+            self.is_sparse,
+            self.units,
+        )
+        return pickled_state[0], pickled_state[1], new_state
+
+    def __setstate__(self, state: Any, **kwargs: Any) -> None:
+        """Sets the state.
+
+        Args:
+            state: The state argument must be a sequence that contains the following
+                   elements version, shape, dtype, isFortan, rawdata.
+            **kwargs: Any additional parameter
+        """
+        # Needed in order to load the object
+        self.number_of_dimensions = state[-3]
+        self.is_sparse = state[-2]
+        self.units = state[-1]
+        super(FeatureArray, self).__setstate__(state[0:-3], **kwargs)
+
+    # pytype: enable=attribute-error
+
+    @staticmethod
+    def _validate_number_of_dimensions(
+        number_of_dimensions: int, input_array: np.ndarray
+    ) -> None:
+        """Validates if the the input array has given number of dimensions.
+
+        Args:
+            number_of_dimensions: number of dimensions
+            input_array: input array
+
+        Raises: ValueError in case the dimensions do not match
+        """
+        _sub_array = input_array
+        dim = 0
+        # Go number_of_dimensions into the given input_array
+        for i in range(1, number_of_dimensions + 1):
+            _sub_array = _sub_array[0]
+            if isinstance(_sub_array, scipy.sparse.spmatrix):
+                dim = i
+                break
+            if isinstance(_sub_array, np.ndarray) and _sub_array.shape[0] == 0:
+                # sequence dimension is 0, we are dealing with "fake" features
+                dim = i
+                break
+
+        # If the resulting sub_array is sparse, the remaining number of dimensions
+        # should be at least 2
+        if isinstance(_sub_array, scipy.sparse.spmatrix):
+            if dim > 2:
+                raise ValueError(
+                    f"Given number of dimensions '{number_of_dimensions}' does not "
+                    f"match dimensions of given input array: {input_array}."
+                )
+        elif isinstance(_sub_array, np.ndarray) and _sub_array.shape[0] == 0:
+            # sequence dimension is 0, we are dealing with "fake" features,
+            # but they should be of dim 2
+            if dim > 2:
+                raise ValueError(
+                    f"Given number of dimensions '{number_of_dimensions}' does not "
+                    f"match dimensions of given input array: {input_array}."
+                )
+        # If the resulting sub_array is dense, the sub_array should be a single number
+        elif not np.issubdtype(type(_sub_array), np.integer) and not isinstance(
+            _sub_array, (np.float32, np.float64)
+        ):
+            raise ValueError(
+                f"Given number of dimensions '{number_of_dimensions}' does not match "
+                f"dimensions of given input array: {input_array}."
+            )
+
+
+class FeatureSignature(NamedTuple):
+    """Signature of feature arrays.
+
+    Stores the number of units, the type (sparse vs dense), and the number of
+    dimensions of features.
+    """
+
+    is_sparse: bool
+    units: Optional[int]
+    number_of_dimensions: int
+
+
+# Mapping of attribute name and feature name to a list of feature arrays representing
+# the actual features
+# For example:
+# "text" -> { "sentence": [
+#   "feature array containing dense features for every training example",
+#   "feature array containing sparse features for every training example"
+# ]}
+Data = Dict[Text, Dict[Text, List[FeatureArray]]]
+
+
+class RasaModelData:
+    """Data object used for all RasaModels.
+
+    It contains all features needed to train the models.
+    'data' is a mapping of attribute name, e.g. TEXT, INTENT, etc., and feature name,
+    e.g. SENTENCE, SEQUENCE, etc., to a list of feature arrays representing the actual
+    features.
+    'label_key' and 'label_sub_key' point to the labels inside 'data'. For
+    example, if your intent labels are stored under INTENT -> IDS, 'label_key' would
+    be "INTENT" and 'label_sub_key' would be "IDS".
+    """
+
+    def __init__(
+        self,
+        label_key: Optional[Text] = None,
+        label_sub_key: Optional[Text] = None,
+        data: Optional[Data] = None,
+    ) -> None:
+        """Initializes the RasaModelData object.
+
+        Args:
+            label_key: the key of a label used for balancing, etc.
+            label_sub_key: the sub key of a label used for balancing, etc.
+            data: the data holding the features
+        """
+        self.data = data or defaultdict(lambda: defaultdict(list))
+        self.label_key = label_key
+        self.label_sub_key = label_sub_key
+        # should be updated when features are added
+        self.num_examples = self.number_of_examples()
+        self.sparse_feature_sizes: Dict[Text, Dict[Text, List[int]]] = {}
+
+    @overload
+    def get(self, key: Text, sub_key: Text) -> List[FeatureArray]:
+        ...
+
+    @overload
+    def get(self, key: Text, sub_key: None = ...) -> Dict[Text, List[FeatureArray]]:
+        ...
+
+    def get(
+        self, key: Text, sub_key: Optional[Text] = None
+    ) -> Union[Dict[Text, List[FeatureArray]], List[FeatureArray]]:
+        """Get the data under the given keys.
+
+        Args:
+            key: The key.
+            sub_key: The optional sub key.
+
+        Returns:
+            The requested data.
+        """
+        if sub_key is None and key in self.data:
+            return self.data[key]
+
+        if sub_key and key in self.data and sub_key in self.data[key]:
+            return self.data[key][sub_key]
+
+        return []
+
+    def items(self) -> ItemsView:
+        """Return the items of the data attribute.
+
+        Returns:
+            The items of data.
+        """
+        return self.data.items()
+
+    def values(self) -> Any:
+        """Return the values of the data attribute.
+
+        Returns:
+            The values of data.
+        """
+        return self.data.values()
+
+    def keys(self, key: Optional[Text] = None) -> List[Text]:
+        """Return the keys of the data attribute.
+
+        Args:
+            key: The optional key.
+
+        Returns:
+            The keys of the data.
+        """
+        if key is None:
+            return list(self.data.keys())
+
+        if key in self.data:
+            return list(self.data[key].keys())
+
+        return []
+
+    def sort(self) -> None:
+        """Sorts data according to its keys."""
+        for key, attribute_data in self.data.items():
+            self.data[key] = OrderedDict(sorted(attribute_data.items()))
+        self.data = OrderedDict(sorted(self.data.items()))
+
+    def first_data_example(self) -> Data:
+        """Return the data with just one feature example per key, sub-key.
+
+        Returns:
+            The simplified data.
+        """
+        out_data: Data = {}
+        for key, attribute_data in self.data.items():
+            out_data[key] = {}
+            for sub_key, features in attribute_data.items():
+                feature_slices = [feature[:1] for feature in features]
+                out_data[key][sub_key] = cast(List[FeatureArray], feature_slices)
+        return out_data
+
+    def does_feature_exist(self, key: Text, sub_key: Optional[Text] = None) -> bool:
+        """Check if feature key (and sub-key) is present and features are available.
+
+        Args:
+            key: The key.
+            sub_key: The optional sub-key.
+
+        Returns:
+            False, if no features for the given keys exists, True otherwise.
+        """
+        return not self.does_feature_not_exist(key, sub_key)
+
+    def does_feature_not_exist(self, key: Text, sub_key: Optional[Text] = None) -> bool:
+        """Check if feature key (and sub-key) is present and features are available.
+
+        Args:
+            key: The key.
+            sub_key: The optional sub-key.
+
+        Returns:
+            True, if no features for the given keys exists, False otherwise.
+        """
+        if sub_key:
+            return (
+                key not in self.data
+                or not self.data[key]
+                or sub_key not in self.data[key]
+                or not self.data[key][sub_key]
+            )
+
+        return key not in self.data or not self.data[key]
+
+    def is_empty(self) -> bool:
+        """Checks if data is set."""
+        return not self.data
+
+    def number_of_examples(self, data: Optional[Data] = None) -> int:
+        """Obtain number of examples in data.
+
+        Args:
+            data: The data.
+
+        Raises: A ValueError if number of examples differ for different features.
+
+        Returns:
+            The number of examples in data.
+        """
+        if not data:
+            data = self.data
+
+        if not data:
+            return 0
+
+        example_lengths = [
+            len(f)
+            for attribute_data in data.values()
+            for features in attribute_data.values()
+            for f in features
+        ]
+
+        if not example_lengths:
+            return 0
+
+        # check if number of examples is the same for all values
+        if not all(length == example_lengths[0] for length in example_lengths):
+            raise ValueError(
+                f"Number of examples differs for keys '{data.keys()}'. Number of "
+                f"examples should be the same for all data."
+            )
+
+        return example_lengths[0]
+
+    def number_of_units(self, key: Text, sub_key: Text) -> int:
+        """Get the number of units of the given key.
+
+        Args:
+            key: The key.
+            sub_key: The optional sub-key.
+
+        Returns:
+            The number of units.
+        """
+        if key not in self.data or sub_key not in self.data[key]:
+            return 0
+
+        units = 0
+        for features in self.data[key][sub_key]:
+            if len(features) > 0:
+                units += features.units  # type: ignore[operator]
+
+        return units
+
+    def add_data(self, data: Data, key_prefix: Optional[Text] = None) -> None:
+        """Add incoming data to data.
+
+        Args:
+            data: The data to add.
+            key_prefix: Optional key prefix to use in front of the key value.
+        """
+        for key, attribute_data in data.items():
+            for sub_key, features in attribute_data.items():
+                if key_prefix:
+                    self.add_features(f"{key_prefix}{key}", sub_key, features)
+                else:
+                    self.add_features(key, sub_key, features)
+
+    def update_key(
+        self, from_key: Text, from_sub_key: Text, to_key: Text, to_sub_key: Text
+    ) -> None:
+        """Copies the features under the given keys to the new keys and deletes the old.
+
+        Args:
+            from_key: current feature key
+            from_sub_key: current feature sub-key
+            to_key: new key for feature
+            to_sub_key: new sub-key for feature
+        """
+        if from_key not in self.data or from_sub_key not in self.data[from_key]:
+            return
+
+        if to_key not in self.data:
+            self.data[to_key] = {}
+        self.data[to_key][to_sub_key] = self.get(from_key, from_sub_key)
+        del self.data[from_key][from_sub_key]
+
+        if not self.data[from_key]:
+            del self.data[from_key]
+
+    def add_features(
+        self, key: Text, sub_key: Text, features: Optional[List[FeatureArray]]
+    ) -> None:
+        """Add list of features to data under specified key.
+
+        Should update number of examples.
+
+        Args:
+            key: The key
+            sub_key: The sub-key
+            features: The features to add.
+        """
+        if features is None:
+            return
+
+        for feature_array in features:
+            if len(feature_array) > 0:
+                self.data[key][sub_key].append(feature_array)
+
+        if not self.data[key][sub_key]:
+            del self.data[key][sub_key]
+
+        # update number of examples
+        self.num_examples = self.number_of_examples()
+
+    def add_lengths(
+        self, key: Text, sub_key: Text, from_key: Text, from_sub_key: Text
+    ) -> None:
+        """Adds a feature array of lengths of sequences to data under given key.
+
+        Args:
+            key: The key to add the lengths to
+            sub_key: The sub-key to add the lengths to
+            from_key: The key to take the lengths from
+            from_sub_key: The sub-key to take the lengths from
+        """
+        if not self.data.get(from_key) or not self.data.get(from_key, {}).get(
+            from_sub_key
+        ):
+            return
+
+        self.data[key][sub_key] = []
+
+        for features in self.data[from_key][from_sub_key]:
+            if len(features) == 0:
+                continue
+
+            if features.number_of_dimensions == 4:
+                lengths = FeatureArray(
+                    ragged_array_to_ndarray(
+                        [
+                            # add one more dim so that dialogue dim
+                            # would be a sequence
+                            np.array([[[x.shape[0]]] for x in _features])
+                            for _features in features
+                        ]
+                    ),
+                    number_of_dimensions=4,
+                )
+            else:
+                lengths = FeatureArray(
+                    np.array([x.shape[0] for x in features]), number_of_dimensions=1
+                )
+            self.data[key][sub_key].extend([lengths])
+            break
+
+    def add_sparse_feature_sizes(
+        self, sparse_feature_sizes: Dict[Text, Dict[Text, List[int]]]
+    ) -> None:
+        """Adds a dictionary of feature sizes for different attributes.
+
+        Args:
+            sparse_feature_sizes: a dictionary of attribute that has sparse
+                           features to a dictionary of a feature type
+                           to a list of different sparse feature sizes.
+        """
+        self.sparse_feature_sizes = sparse_feature_sizes
+
+    def get_sparse_feature_sizes(self) -> Dict[Text, Dict[Text, List[int]]]:
+        """Get feature sizes of the model.
+
+        sparse_feature_sizes is a dictionary of attribute that has sparse features to
+        a dictionary of a feature type to a list of different sparse feature sizes.
+
+        Returns:
+            A dictionary of key and sub-key to a list of feature signatures
+            (same structure as the data attribute).
+        """
+        return self.sparse_feature_sizes
+
+    def split(
+        self, number_of_test_examples: int, random_seed: int
+    ) -> Tuple["RasaModelData", "RasaModelData"]:
+        """Create random hold out test set using stratified split.
+
+        Args:
+            number_of_test_examples: Number of test examples.
+            random_seed: Random seed.
+
+        Returns:
+            A tuple of train and test RasaModelData.
+        """
+        self._check_label_key()
+
+        if self.label_key is None or self.label_sub_key is None:
+            # randomly split data as no label key is set
+            multi_values = [
+                v
+                for attribute_data in self.data.values()
+                for data in attribute_data.values()
+                for v in data
+            ]
+            solo_values: List[Any] = [
+                []
+                for attribute_data in self.data.values()
+                for data in attribute_data.values()
+                for _ in data
+            ]
+            stratify = None
+        else:
+            # make sure that examples for each label value are in both split sets
+            label_ids = self._create_label_ids(
+                self.data[self.label_key][self.label_sub_key][0]
+            )
+            label_counts: Dict[int, int] = dict(
+                zip(
+                    *np.unique(
+                        label_ids,
+                        return_counts=True,
+                        axis=0,
+                    )
+                )
+            )
+
+            self._check_train_test_sizes(number_of_test_examples, label_counts)
+
+            counts = np.array([label_counts[label] for label in label_ids])
+            # we perform stratified train test split,
+            # which insures every label is present in the train and test data
+            # this operation can be performed only for labels
+            # that contain several data points
+            multi_values = [
+                f[counts > 1].view(FeatureArray)
+                for attribute_data in self.data.values()
+                for features in attribute_data.values()
+                for f in features
+            ]
+            # collect data points that are unique for their label
+            solo_values = [
+                f[counts == 1]
+                for attribute_data in self.data.values()
+                for features in attribute_data.values()
+                for f in features
+            ]
+
+            stratify = label_ids[counts > 1]
+
+        output_values = train_test_split(
+            *multi_values,
+            test_size=number_of_test_examples,
+            random_state=random_seed,
+            stratify=stratify,
+        )
+
+        return self._convert_train_test_split(output_values, solo_values)
+
+    def get_signature(
+        self, data: Optional[Data] = None
+    ) -> Dict[Text, Dict[Text, List[FeatureSignature]]]:
+        """Get signature of RasaModelData.
+
+        Signature stores the shape and whether features are sparse or not for every key.
+
+        Returns:
+            A dictionary of key and sub-key to a list of feature signatures
+            (same structure as the data attribute).
+        """
+        if not data:
+            data = self.data
+
+        return {
+            key: {
+                sub_key: [
+                    FeatureSignature(f.is_sparse, f.units, f.number_of_dimensions)
+                    for f in features
+                ]
+                for sub_key, features in attribute_data.items()
+            }
+            for key, attribute_data in data.items()
+        }
+
+    def shuffled_data(self, data: Data) -> Data:
+        """Shuffle model data.
+
+        Args:
+            data: The data to shuffle
+
+        Returns:
+            The shuffled data.
+        """
+        ids = np.random.permutation(self.num_examples)
+        return self._data_for_ids(data, ids)
+
+    def balanced_data(self, data: Data, batch_size: int, shuffle: bool) -> Data:
+        """Mix model data to account for class imbalance.
+
+        This batching strategy puts rare classes approximately in every other batch,
+        by repeating them. Mimics stratified batching, but also takes into account
+        that more populated classes should appear more often.
+
+        Args:
+            data: The data.
+            batch_size: The batch size.
+            shuffle: Boolean indicating whether to shuffle the data or not.
+
+        Returns:
+            The balanced data.
+        """
+        self._check_label_key()
+
+        # skip balancing if labels are token based
+        if (
+            self.label_key is None
+            or self.label_sub_key is None
+            or data[self.label_key][self.label_sub_key][0][0].size > 1
+        ):
+            return data
+
+        label_ids = self._create_label_ids(data[self.label_key][self.label_sub_key][0])
+
+        unique_label_ids, counts_label_ids = np.unique(
+            label_ids, return_counts=True, axis=0
+        )
+        num_label_ids = len(unique_label_ids)
+
+        # group data points by their label
+        # need to call every time, so that the data is shuffled inside each class
+        data_by_label = self._split_by_label_ids(data, label_ids, unique_label_ids)
+
+        # running index inside each data grouped by labels
+        data_idx = [0] * num_label_ids
+        # number of cycles each label was passed
+        num_data_cycles = [0] * num_label_ids
+        # if a label was skipped in current batch
+        skipped = [False] * num_label_ids
+
+        new_data: DefaultDict[
+            Text, DefaultDict[Text, List[List[FeatureArray]]]
+        ] = defaultdict(lambda: defaultdict(list))
+
+        while min(num_data_cycles) == 0:
+            if shuffle:
+                indices_of_labels = np.random.permutation(num_label_ids)
+            else:
+                indices_of_labels = np.asarray(range(num_label_ids))
+
+            for index in indices_of_labels:
+                if num_data_cycles[index] > 0 and not skipped[index]:
+                    skipped[index] = True
+                    continue
+
+                skipped[index] = False
+
+                index_batch_size = (
+                    int(counts_label_ids[index] / self.num_examples * batch_size) + 1
+                )
+
+                for key, attribute_data in data_by_label[index].items():
+                    for sub_key, features in attribute_data.items():
+                        for i, f in enumerate(features):
+                            if len(new_data[key][sub_key]) < i + 1:
+                                new_data[key][sub_key].append([])
+                            new_data[key][sub_key][i].append(
+                                f[data_idx[index] : data_idx[index] + index_batch_size]
+                            )
+
+                data_idx[index] += index_batch_size
+                if data_idx[index] >= counts_label_ids[index]:
+                    num_data_cycles[index] += 1
+                    data_idx[index] = 0
+
+                if min(num_data_cycles) > 0:
+                    break
+
+        final_data: Data = defaultdict(lambda: defaultdict(list))
+        for key, attribute_data in new_data.items():
+            for sub_key, features in attribute_data.items():
+                for f in features:
+                    final_data[key][sub_key].append(
+                        FeatureArray(
+                            np.concatenate(f),
+                            number_of_dimensions=f[0].number_of_dimensions,
+                        )
+                    )
+
+        return final_data
+
+    def _check_train_test_sizes(
+        self, number_of_test_examples: int, label_counts: Dict[Any, int]
+    ) -> None:
+        """Check whether the test data set is too large or too small.
+
+        Args:
+            number_of_test_examples: number of test examples
+            label_counts: number of labels
+
+        Raises:
+            A ValueError if the number of examples does not fit.
+        """
+        if number_of_test_examples >= self.num_examples - len(label_counts):
+            raise ValueError(
+                f"Test set of {number_of_test_examples} is too large. Remaining "
+                f"train set should be at least equal to number of classes "
+                f"{len(label_counts)}."
+            )
+        if number_of_test_examples < len(label_counts):
+            raise ValueError(
+                f"Test set of {number_of_test_examples} is too small. It should "
+                f"be at least equal to number of classes {label_counts}."
+            )
+
+    @staticmethod
+    def _data_for_ids(data: Optional[Data], ids: np.ndarray) -> Data:
+        """Filter model data by ids.
+
+        Args:
+            data: The data to filter
+            ids: The ids
+
+        Returns:
+            The filtered data
+        """
+        new_data: Data = defaultdict(lambda: defaultdict(list))
+
+        if data is None:
+            return new_data
+
+        for key, attribute_data in data.items():
+            for sub_key, features in attribute_data.items():
+                for f in features:
+                    new_data[key][sub_key].append(f[ids])
+        return new_data
+
+    def _split_by_label_ids(
+        self, data: Optional[Data], label_ids: np.ndarray, unique_label_ids: np.ndarray
+    ) -> List["RasaModelData"]:
+        """Reorganize model data into a list of model data with the same labels.
+
+        Args:
+            data: The data
+            label_ids: The label ids
+            unique_label_ids: The unique label ids
+
+        Returns:
+            Reorganized RasaModelData
+        """
+        label_data = []
+        for label_id in unique_label_ids:
+            matching_ids = np.array(label_ids) == label_id
+            label_data.append(
+                RasaModelData(
+                    self.label_key,
+                    self.label_sub_key,
+                    self._data_for_ids(data, matching_ids),
+                )
+            )
+        return label_data
+
+    def _check_label_key(self) -> None:
+        """Check if the label key exists.
+
+        Raises:
+            ValueError if the label key and sub-key is not in data.
+        """
+        if (
+            self.label_key is not None
+            and self.label_sub_key is not None
+            and (
+                self.label_key not in self.data
+                or self.label_sub_key not in self.data[self.label_key]
+                or len(self.data[self.label_key][self.label_sub_key]) > 1
+            )
+        ):
+            raise ValueError(
+                f"Key '{self.label_key}.{self.label_sub_key}' not in RasaModelData."
+            )
+
+    def _convert_train_test_split(
+        self, output_values: List[Any], solo_values: List[Any]
+    ) -> Tuple["RasaModelData", "RasaModelData"]:
+        """Converts the output of sklearn's train_test_split into model data.
+
+        Args:
+            output_values: output values of sklearn's train_test_split
+            solo_values: list of solo values
+
+        Returns:
+            The test and train RasaModelData
+        """
+        data_train: DefaultDict[
+            Text, DefaultDict[Text, List[FeatureArray]]
+        ] = defaultdict(lambda: defaultdict(list))
+        data_val: DefaultDict[Text, DefaultDict[Text, List[Any]]] = defaultdict(
+            lambda: defaultdict(list)
+        )
+
+        # output_values = x_train, x_val, y_train, y_val, z_train, z_val, etc.
+        # order is kept, e.g. same order as model data keys
+
+        # train datasets have an even index
+        index = 0
+        for key, attribute_data in self.data.items():
+            for sub_key, features in attribute_data.items():
+                for f in features:
+                    data_train[key][sub_key].append(
+                        self._combine_features(
+                            output_values[index * 2],
+                            solo_values[index],
+                            f.number_of_dimensions,
+                        )
+                    )
+                    index += 1
+
+        # val datasets have an odd index
+        index = 0
+        for key, attribute_data in self.data.items():
+            for sub_key, features in attribute_data.items():
+                for _ in features:
+                    data_val[key][sub_key].append(output_values[(index * 2) + 1])
+                    index += 1
+
+        return (
+            RasaModelData(self.label_key, self.label_sub_key, data_train),
+            RasaModelData(self.label_key, self.label_sub_key, data_val),
+        )
+
+    @staticmethod
+    def _combine_features(
+        feature_1: Union[np.ndarray, scipy.sparse.spmatrix],
+        feature_2: Union[np.ndarray, scipy.sparse.spmatrix],
+        number_of_dimensions: Optional[int] = 1,
+    ) -> FeatureArray:
+        """Concatenate features.
+
+        Args:
+            feature_1: Features to concatenate.
+            feature_2: Features to concatenate.
+
+        Returns:
+            The combined features.
+        """
+        if isinstance(feature_1, scipy.sparse.spmatrix) and isinstance(
+            feature_2, scipy.sparse.spmatrix
+        ):
+            if feature_2.shape[0] == 0:
+                return FeatureArray(feature_1, number_of_dimensions)
+            if feature_1.shape[0] == 0:
+                return FeatureArray(feature_2, number_of_dimensions)
+            return FeatureArray(
+                scipy.sparse.vstack([feature_1, feature_2]), number_of_dimensions
+            )
+        return FeatureArray(
+            np.concatenate([feature_1, feature_2]),
+            number_of_dimensions,
+        )
+
+    @staticmethod
+    def _create_label_ids(label_ids: FeatureArray) -> np.ndarray:
+        """Convert various size label_ids into single dim array.
+
+        For multi-label y, map each distinct row to a string representation
+        using join because str(row) uses an ellipsis if len(row) > 1000.
+        Idea taken from sklearn's stratify split.
+
+        Args:
+            label_ids: The label ids.
+
+        Raises:
+            ValueError if dimensionality of label ids is not supported
+
+        Returns:
+            The single dim label array.
+        """
+        if label_ids.ndim == 1:
+            return label_ids
+
+        if label_ids.ndim == 2 and label_ids.shape[-1] == 1:
+            return label_ids[:, 0]
+
+        if label_ids.ndim == 2:
+            return np.array([" ".join(row.astype("str")) for row in label_ids])
+
+        if label_ids.ndim == 3 and label_ids.shape[-1] == 1:
+            return np.array([" ".join(row.astype("str")) for row in label_ids[:, :, 0]])
+
+        raise ValueError("Unsupported label_ids dimensions")