rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/utils/tensorflow/rasa_layers.py

@@ -0,0 +1,1094 @@
+import tensorflow as tf
+import numpy as np
+from typing import Text, List, Dict, Any, Union, Optional, Tuple, Callable
+
+from rasa.shared.nlu.constants import TEXT
+from rasa.utils.tensorflow.model_data import FeatureSignature
+from rasa.utils.tensorflow.constants import (
+    REGULARIZATION_CONSTANT,
+    CONNECTION_DENSITY,
+    NUM_TRANSFORMER_LAYERS,
+    TRANSFORMER_SIZE,
+    NUM_HEADS,
+    UNIDIRECTIONAL_ENCODER,
+    KEY_RELATIVE_ATTENTION,
+    VALUE_RELATIVE_ATTENTION,
+    MAX_RELATIVE_POSITION,
+    MASKED_LM,
+    HIDDEN_LAYERS_SIZES,
+    DROP_RATE,
+    SPARSE_INPUT_DROPOUT,
+    DENSE_INPUT_DROPOUT,
+    DENSE_DIMENSION,
+    CONCAT_DIMENSION,
+    DROP_RATE_ATTENTION,
+    SEQUENCE,
+    SENTENCE,
+)
+from rasa.utils.tensorflow import layers
+from rasa.utils.tensorflow.exceptions import TFLayerConfigException
+from rasa.utils.tensorflow.transformer import TransformerEncoder
+from rasa.nlu.constants import DEFAULT_TRANSFORMER_SIZE
+
+
+class RasaCustomLayer(tf.keras.layers.Layer):
+    """Parent class for all classes in `rasa_layers.py`.
+
+    Allows a shared implementation for adjusting `DenseForSparse`
+    layers during incremental training.
+
+    During fine-tuning, sparse feature sizes might change due to addition of new data.
+    If this happens, we need to adjust our `DenseForSparse` layers to a new size.
+    `ConcatenateSparseDenseFeatures`, `RasaSequenceLayer` and
+    `RasaFeatureCombiningLayer` all inherit from `RasaCustomLayer` and thus can
+    change their own `DenseForSparse` layers if it's needed.
+    """
+
+    def adjust_sparse_layers_for_incremental_training(
+        self,
+        new_sparse_feature_sizes: Dict[Text, Dict[Text, List[int]]],
+        old_sparse_feature_sizes: Dict[Text, Dict[Text, List[int]]],
+        reg_lambda: float,
+    ) -> None:
+        """Finds and adjusts `DenseForSparse` layers during incremental training.
+
+        Recursively looks through the layers until it finds all the `DenseForSparse`
+        ones and adjusts those which have their sparse feature sizes increased.
+
+        This function heavily relies on the name of `DenseForSparse` layer being
+        in the following format - f"sparse_to_dense.{attribute}_{feature_type}" -
+        in order to correctly extract the attribute and feature type.
+
+        New and old sparse feature sizes could look like this:
+        {TEXT: {FEATURE_TYPE_SEQUENCE: [4, 24, 128], FEATURE_TYPE_SENTENCE: [4, 128]}}
+
+        Args:
+            new_sparse_feature_sizes: sizes of current sparse features.
+            old_sparse_feature_sizes: sizes of sparse features the model was
+                                      previously trained on.
+            reg_lambda: regularization constant.
+        """
+        for name, layer in self._tf_layers.items():
+            if isinstance(layer, RasaCustomLayer):
+                layer.adjust_sparse_layers_for_incremental_training(
+                    new_sparse_feature_sizes=new_sparse_feature_sizes,
+                    old_sparse_feature_sizes=old_sparse_feature_sizes,
+                    reg_lambda=reg_lambda,
+                )
+            elif isinstance(layer, layers.DenseForSparse):
+                attribute = layer.get_attribute()
+                feature_type = layer.get_feature_type()
+                if (
+                    attribute in new_sparse_feature_sizes
+                    and feature_type in new_sparse_feature_sizes[attribute]
+                ):
+                    new_feature_sizes = new_sparse_feature_sizes[attribute][
+                        feature_type
+                    ]
+                    old_feature_sizes = old_sparse_feature_sizes[attribute][
+                        feature_type
+                    ]
+                    if sum(new_feature_sizes) > sum(old_feature_sizes):
+                        self._tf_layers[name] = self._replace_dense_for_sparse_layer(
+                            layer_to_replace=layer,
+                            new_sparse_feature_sizes=new_feature_sizes,
+                            old_sparse_feature_sizes=old_feature_sizes,
+                            attribute=attribute,
+                            feature_type=feature_type,
+                            reg_lambda=reg_lambda,
+                        )
+
+    @staticmethod
+    def _replace_dense_for_sparse_layer(
+        layer_to_replace: layers.DenseForSparse,
+        new_sparse_feature_sizes: List[int],
+        old_sparse_feature_sizes: List[int],
+        attribute: Text,
+        feature_type: Text,
+        reg_lambda: float,
+    ) -> layers.DenseForSparse:
+        """Replaces a `DenseForSparse` layer with a new one.
+
+        Replaces an existing `DenseForSparse` layer with a new one
+        in order to adapt it to incremental training.
+
+        Args:
+            layer_to_replace: a `DenseForSparse` layer that is used to create a new one.
+            new_sparse_feature_sizes: sizes of sparse features that will be
+                                      the input of the layer.
+            old_sparse_feature_sizes: sizes of sparse features that used to be
+                                      the input of the layer.
+            attribute: an attribute of the data fed to the layer.
+            feature_type: a feature type of the data fed to the layer.
+            reg_lambda: regularization constant.
+
+        Returns:
+            New `DenseForSparse` layer.
+        """
+        kernel = layer_to_replace.get_kernel().numpy()
+        bias = layer_to_replace.get_bias()
+        if bias is not None:
+            bias = bias.numpy()
+        units = layer_to_replace.get_units()
+        # split kernel by feature sizes to update the layer accordingly
+        kernel_splits = []
+        splitting_index = 0
+        for size in old_sparse_feature_sizes:
+            kernel_splits.append(kernel[splitting_index : splitting_index + size, :])
+            splitting_index += size
+        additional_sizes = [
+            new_size - old_size
+            for new_size, old_size in zip(
+                new_sparse_feature_sizes, old_sparse_feature_sizes
+            )
+        ]
+        std, mean = np.std(kernel), np.mean(kernel)
+        additional_weights = [
+            np.random.normal(mean, std, size=(num_rows, units)).astype(np.float32)
+            for num_rows in additional_sizes
+        ]
+        merged_weights = [
+            np.vstack((existing, new))
+            for existing, new in zip(kernel_splits, additional_weights)
+        ]
+        # stack each merged weight to form a new weight tensor
+        new_weights = np.vstack(merged_weights)
+        kernel_init = tf.constant_initializer(new_weights)
+        bias_init = tf.constant_initializer(bias) if bias is not None else None
+        new_layer = layers.DenseForSparse(
+            name=f"sparse_to_dense.{attribute}_{feature_type}",
+            reg_lambda=reg_lambda,
+            units=units,
+            use_bias=bias is not None,
+            kernel_initializer=kernel_init,
+            bias_initializer=bias_init,
+        )
+        return new_layer
+
+
+class ConcatenateSparseDenseFeatures(RasaCustomLayer):
+    """Combines multiple sparse and dense feature tensors into one dense tensor.
+
+    This layer combines features from various featurisers into a single feature array
+    per input example. All features must be of the same feature type, i.e. sentence-
+    level or sequence-level (token-level).
+
+    The layer combines a given list of tensors (whether sparse or dense) by:
+    1. converting sparse tensors into dense ones
+    2. optionally, applying dropout to sparse tensors before and/or after the conversion
+    3. concatenating all tensors along the last dimension
+
+    Arguments:
+        attribute: Name of attribute (e.g. `text` or `label`) whose features will be
+            processed.
+        feature_type: Feature type to be processed -- `sequence` or `sentence`.
+        feature_type_signature: A list of signatures for the given attribute and feature
+            type.
+        config: A model config for correctly parametrising the layer.
+
+    Input shape:
+        Tuple containing one list of N-D tensors, each with shape: `(batch_size, ...,
+        input_dim)`.
+        All dense tensors must have the same shape, except possibly the last dimension.
+        All sparse tensors must have the same shape, including the last dimension.
+
+    Output shape:
+        N-D tensor with shape: `(batch_size, ..., units)` where `units` is the sum of
+        the last dimension sizes across all input tensors, with sparse tensors instead
+        contributing `config[DENSE_DIMENSION][attribute]` units each.
+
+    Raises:
+        A `TFLayerConfigException` if no feature signatures are provided.
+
+    Attributes:
+        output_units: The last dimension size of the layer's output.
+    """
+
+    SPARSE_DROPOUT = "sparse_dropout"
+    SPARSE_TO_DENSE = "sparse_to_dense"
+    DENSE_DROPOUT = "dense_dropout"
+
+    def __init__(
+        self,
+        attribute: Text,
+        feature_type: Text,
+        feature_type_signature: List[FeatureSignature],
+        config: Dict[Text, Any],
+    ) -> None:
+        """Creates a new `ConcatenateSparseDenseFeatures` object."""
+        if not feature_type_signature:
+            raise TFLayerConfigException(
+                "The feature type signature must contain some feature signatures."
+            )
+
+        super().__init__(
+            name=f"concatenate_sparse_dense_features_{attribute}_{feature_type}"
+        )
+
+        self._check_sparse_input_units(feature_type_signature)
+
+        self.output_units = self._calculate_output_units(
+            attribute, feature_type_signature, config
+        )
+
+        # Prepare dropout and sparse-to-dense layers if any sparse tensors are expected
+        self._tf_layers: Dict[Text, tf.keras.layers.Layer] = {}
+        if any([signature.is_sparse for signature in feature_type_signature]):
+            self._prepare_layers_for_sparse_tensors(attribute, feature_type, config)
+
+    def _check_sparse_input_units(
+        self, feature_type_signature: List[FeatureSignature]
+    ) -> None:
+        """Checks that all sparse features have the same last dimension size."""
+        sparse_units = [
+            feature_sig.units
+            for feature_sig in feature_type_signature
+            if feature_sig.is_sparse
+        ]
+        if len(set(sparse_units)) > 1:
+            raise TFLayerConfigException(
+                f"All sparse features must have the same last dimension size but found "
+                f"different sizes: {set(sparse_units)}."
+            )
+
+    def _prepare_layers_for_sparse_tensors(
+        self, attribute: Text, feature_type: Text, config: Dict[Text, Any]
+    ) -> None:
+        """Sets up sparse tensor pre-processing before combining with dense ones."""
+        # For optionally applying dropout to sparse tensors
+        if config[SPARSE_INPUT_DROPOUT]:
+            self._tf_layers[self.SPARSE_DROPOUT] = layers.SparseDropout(
+                rate=config[DROP_RATE]
+            )
+
+        # For converting sparse tensors to dense
+        self._tf_layers[self.SPARSE_TO_DENSE] = layers.DenseForSparse(
+            name=f"sparse_to_dense.{attribute}_{feature_type}",
+            units=config[DENSE_DIMENSION][attribute],
+            reg_lambda=config[REGULARIZATION_CONSTANT],
+        )
+
+        # For optionally apply dropout to sparse tensors after they're converted to
+        # dense tensors.
+        if config[DENSE_INPUT_DROPOUT]:
+            self._tf_layers[self.DENSE_DROPOUT] = tf.keras.layers.Dropout(
+                rate=config[DROP_RATE]
+            )
+
+    @staticmethod
+    def _calculate_output_units(
+        attribute: Text,
+        feature_type_signature: List[FeatureSignature],
+        config: Dict[Text, Any],
+    ) -> int:
+        """Determines the output units from the provided feature signatures.
+
+        Sparse features will be turned into dense ones, hence they each contribute with
+        their future dense number of units.
+        """
+        return sum(
+            [
+                config[DENSE_DIMENSION][attribute]
+                if signature.is_sparse
+                else signature.units
+                for signature in feature_type_signature
+            ]
+        )
+
+    def _process_sparse_feature(
+        self, feature: tf.SparseTensor, training: bool
+    ) -> tf.Tensor:
+        """Turns sparse tensor into dense, possibly adds dropout before and/or after."""
+        if self.SPARSE_DROPOUT in self._tf_layers:
+            feature = self._tf_layers[self.SPARSE_DROPOUT](feature, training)
+
+        feature = self._tf_layers[self.SPARSE_TO_DENSE](feature)
+
+        if self.DENSE_DROPOUT in self._tf_layers:
+            feature = self._tf_layers[self.DENSE_DROPOUT](feature, training)
+
+        return feature
+
+    def call(
+        self,
+        inputs: Tuple[List[Union[tf.Tensor, tf.SparseTensor]]],
+        training: bool = False,
+    ) -> tf.Tensor:
+        """Combines sparse and dense feature tensors into one tensor.
+
+        Arguments:
+            inputs: Contains the input tensors, all of the same rank.
+            training: A flag indicating whether the layer should behave in training mode
+                (applying dropout to sparse tensors if applicable) or in inference mode
+                (not applying dropout).
+
+        Returns:
+            Single tensor with all input tensors combined along the last dimension.
+        """
+        features = inputs[0]
+
+        dense_features = []
+        for f in features:
+            if isinstance(f, tf.SparseTensor):
+                f = self._process_sparse_feature(f, training)
+            dense_features.append(f)
+
+        # Now that all features are made dense, concatenate them along the last (units)
+        # dimension.
+        return tf.concat(dense_features, axis=-1)
+
+
+class RasaFeatureCombiningLayer(RasaCustomLayer):
+    """Combines multiple dense or sparse feature tensors into one.
+
+    This layer combines features by following these steps:
+    1. Apply a `ConcatenateSparseDenseFeatures` layer separately to sequence- and
+        sentence-level features, yielding two tensors (one for each feature type).
+    2. Concatenate the sequence- and sentence-level tensors along the sequence dimension
+        by appending sentence-level features at the first available token position after
+        the sequence-level (token-level) features.
+
+    Arguments:
+        attribute: Name of attribute (e.g. `text` or `label`) whose features will be
+            processed.
+        attribute_signature: A dictionary containing two lists of feature signatures,
+            one for each feature type (`sequence` or `sentence`) of the given attribute.
+        config: A model config used for correctly parameterising the layer and the
+            `ConcatenateSparseDenseFeatures` layer it uses internally.
+
+    Input shape:
+        Tuple of three input tensors:
+            sequence_features: List of 3-D dense or sparse tensors, each with shape
+                `(batch_size, max_seq_length, input_dim)` where `input_dim` can be
+                different for sparse vs dense tensors. See the input shape of
+                `ConcatenateSparseDenseFeatures` for more information.
+            sentence_features: List of 3-D dense or sparse tensors, each with shape
+                `(batch_size, 1, input_dim)` where `input_dim` can be different for
+                sparse vs dense tensors, and can differ from that in
+                `sequence_features`. See the input shape of
+                `ConcatenateSparseDenseFeatures` for more information.
+            sequence_feature_lengths: Dense tensor of shape `(batch_size, )`.
+
+    Output shape:
+        combined_features: A 3-D tensor with shape `(batch_size, sequence_length,
+            units)` where `units` is  completely  determined by the internally applied
+            `ConcatenateSparseDenseFeatures` layer and `sequence_length` is the combined
+            length of sequence- and sentence-level features: `max_seq_length + 1` if
+            both feature types are present, `max_seq_length` if only sequence-level
+            features are present, and 1 if only sentence-level features are present).
+        mask_combined_sequence_sentence: A 3-D tensor with shape
+            `(batch_size, sequence_length, 1)`.
+
+    Raises:
+        A `TFLayerConfigException` if no feature signatures are provided.
+
+    Attributes:
+        output_units: The last dimension size of the layer's `combined_features` output.
+    """
+
+    def __init__(
+        self,
+        attribute: Text,
+        attribute_signature: Dict[Text, List[FeatureSignature]],
+        config: Dict[Text, Any],
+    ) -> None:
+        """Creates a new `RasaFeatureCombiningLayer` object."""
+        if not attribute_signature or not (
+            attribute_signature.get(SENTENCE, [])
+            or attribute_signature.get(SEQUENCE, [])
+        ):
+            raise TFLayerConfigException(
+                "The attribute signature must contain some feature signatures."
+            )
+
+        super().__init__(name=f"rasa_feature_combining_layer_{attribute}")
+
+        self._tf_layers: Dict[Text, tf.keras.layers.Layer] = {}
+
+        # Prepare sparse-dense combining layers for each present feature type
+        self._feature_types_present = self._get_present_feature_types(
+            attribute_signature
+        )
+        self._prepare_sparse_dense_concat_layers(attribute, attribute_signature, config)
+
+        # Prepare components for combining sequence- and sentence-level features
+        self._prepare_sequence_sentence_concat(attribute, config)
+
+        self.output_units = self._calculate_output_units(attribute, config)
+
+    @staticmethod
+    def _get_present_feature_types(
+        attribute_signature: Dict[Text, List[FeatureSignature]],
+    ) -> Dict[Text, bool]:
+        """Determines feature types that are present.
+
+        Knowing which feature types are present is important because many downstream
+        operations depend on it, e.g. combining sequence- and sentence-level features
+        is only done if both feature types are present.
+        """
+        return {
+            feature_type: (
+                feature_type in attribute_signature
+                and len(attribute_signature[feature_type]) > 0
+            )
+            for feature_type in [SEQUENCE, SENTENCE]
+        }
+
+    def _prepare_sparse_dense_concat_layers(
+        self,
+        attribute: Text,
+        attribute_signature: Dict[Text, List[FeatureSignature]],
+        config: Dict[Text, Any],
+    ) -> None:
+        """Prepares sparse-dense combining layers for all present feature types."""
+        for feature_type, present in self._feature_types_present.items():
+            if not present:
+                continue
+            self._tf_layers[f"sparse_dense.{feature_type}"] = (
+                ConcatenateSparseDenseFeatures(
+                    attribute=attribute,
+                    feature_type=feature_type,
+                    feature_type_signature=attribute_signature[feature_type],
+                    config=config,
+                )
+            )
+
+    def _prepare_sequence_sentence_concat(
+        self, attribute: Text, config: Dict[Text, Any]
+    ) -> None:
+        """Sets up combining sentence- and sequence-level features (if needed).
+
+        This boils down to preparing for unifying the units of the sequence- and
+        sentence-level features if they differ -- the same number of units is required
+        for combining the features.
+        """
+        if (
+            self._feature_types_present[SEQUENCE]
+            and self._feature_types_present[SENTENCE]
+        ):
+            # The output units of this layer will be based on the output sizes of the
+            # sparse+dense combining layers that are internally applied to all features.
+            sequence_units = self._tf_layers[f"sparse_dense.{SEQUENCE}"].output_units
+            sentence_units = self._tf_layers[f"sparse_dense.{SENTENCE}"].output_units
+
+            # Last dimension needs to be unified if sequence- and sentence-level
+            # features have different sizes, e.g. due to being produced by different
+            # featurizers.
+            if sequence_units != sentence_units:
+                for feature_type in [SEQUENCE, SENTENCE]:
+                    self._tf_layers[
+                        f"unify_dims_before_seq_sent_concat.{feature_type}"
+                    ] = layers.Ffnn(
+                        layer_name_suffix=f"unify_dims.{attribute}_{feature_type}",
+                        layer_sizes=[config[CONCAT_DIMENSION][attribute]],
+                        dropout_rate=config[DROP_RATE],
+                        reg_lambda=config[REGULARIZATION_CONSTANT],
+                        density=config[CONNECTION_DENSITY],
+                    )
+
+    def _calculate_output_units(self, attribute: Text, config: Dict[Text, Any]) -> int:
+        """Calculates the number of output units for this layer class.
+
+        The number depends mainly on whether dimension unification is used or not.
+        """
+        # If dimension unification is used, output units are determined by the unifying
+        # layers.
+        if (
+            f"unify_dims_before_seq_sent_concat.{SEQUENCE}" in self._tf_layers
+            or f"unify_dims_before_seq_sent_concat.{SENTENCE}" in self._tf_layers
+        ):
+            return config[CONCAT_DIMENSION][attribute]
+        # Without dimension unification, the units from the underlying sparse_dense
+        # layers are carried over and should be the same for sequence-level features
+        # (if present) as for sentence-level features.
+        elif self._feature_types_present[SEQUENCE]:
+            return self._tf_layers[f"sparse_dense.{SEQUENCE}"].output_units
+        return self._tf_layers[f"sparse_dense.{SENTENCE}"].output_units
+
+    def _concat_sequence_sentence_features(
+        self,
+        sequence_tensor: tf.Tensor,
+        sentence_tensor: tf.Tensor,
+        mask_combined_sequence_sentence: tf.Tensor,
+    ) -> tf.Tensor:
+        """Concatenates sequence- & sentence-level features along sequence dimension."""
+        # If needed, pass both feature types through a dense layer to bring them to the
+        # same shape.
+        if f"unify_dims_before_seq_sent_concat.{SEQUENCE}" in self._tf_layers:
+            sequence_tensor = self._tf_layers[
+                f"unify_dims_before_seq_sent_concat.{SEQUENCE}"
+            ](sequence_tensor)
+        if f"unify_dims_before_seq_sent_concat.{SENTENCE}" in self._tf_layers:
+            sentence_tensor = self._tf_layers[
+                f"unify_dims_before_seq_sent_concat.{SENTENCE}"
+            ](sentence_tensor)
+
+        # mask_combined_sequence_sentence has for each input example a sequence of 1s of
+        # the length seq_length+1, where seq_length is the number of real tokens. The
+        # rest is 0s which form a padding up to the max. sequence length + 1 (max.
+        # number of real tokens + 1). Here the mask is turned into a mask that has 0s
+        # everywhere and 1 only at the immediate next position after the last real
+        # token's position for a given input example. Example (batch size = 2, sequence
+        # lengths = [1, 2]):
+        # [[[1], [0], [0]],     ___\   [[[0], [1], [0]],
+        #  [[1], [1], [0]]]        /    [[0], [0], [1]]]
+        sentence_feature_positions_mask = (
+            mask_combined_sequence_sentence
+            * tf.math.cumprod(
+                1 - mask_combined_sequence_sentence,
+                axis=1,
+                exclusive=True,
+                reverse=True,
+            )
+        )
+
+        # The new mask is used to distribute the sentence features at the sequence
+        # positions marked by 1s. The sentence features' dimensionality effectively
+        # changes from `(batch_size, 1, feature_dim)` to `(batch_size, max_seq_length+1,
+        # feature_dim)`, but the array is sparse, with real features present only at
+        # positions determined by 1s in the mask.
+        sentence_tensor = sentence_feature_positions_mask * sentence_tensor
+
+        # Padding of sequence-level features is increased by 1 in the sequence
+        # dimension to match the shape of modified sentence-level features.
+        sequence_tensor = tf.pad(sequence_tensor, [[0, 0], [0, 1], [0, 0]])
+
+        # Sequence- and sentence-level features effectively get concatenated by
+        # summing the two padded feature arrays like this (batch size = 1):
+        # [[seq1, seq2, seq3, 0, 0]] + [[0, 0, 0, sent1, 0]] =
+        # = [[seq1, seq2, seq3, sent1, 0]]
+        return sequence_tensor + sentence_tensor
+
+    def _combine_sequence_level_features(
+        self,
+        sequence_features: List[Union[tf.Tensor, tf.SparseTensor]],
+        mask_sequence: tf.Tensor,
+        training: bool,
+    ) -> Optional[tf.Tensor]:
+        """Processes & combines sequence-level features if any are present."""
+        if self._feature_types_present[SEQUENCE]:
+            sequence_features_combined = self._tf_layers[f"sparse_dense.{SEQUENCE}"](
+                (sequence_features,), training=training
+            )
+
+            # Apply mask which has 1s at positions of real tokens and 0s at all padded
+            # token positions. This is needed because the sparse+dense combining layer
+            # might've turned some fake (padded) features (i.e. 0s) into non-zero
+            # numbers and we want those to become zeros again.
+            # This step isn't needed for sentence-level features because those are never
+            # padded -- the effective sequence length in their case is always 1.
+            return sequence_features_combined * mask_sequence
+
+        return None
+
+    def _combine_sentence_level_features(
+        self,
+        sentence_features: List[Union[tf.Tensor, tf.SparseTensor]],
+        sequence_feature_lengths: tf.Tensor,
+        training: bool,
+    ) -> Tuple[Optional[tf.Tensor], Optional[tf.Tensor]]:
+        """Processes & combines sentence-level features if any are present."""
+        if self._feature_types_present[SENTENCE]:
+            sentence_features_combined = self._tf_layers[f"sparse_dense.{SENTENCE}"](
+                (sentence_features,), training=training
+            )
+            # Sentence-level features have sequence dimension of length 1, add it to
+            # sequence-level feature lengths.
+            combined_sequence_sentence_feature_lengths = sequence_feature_lengths + 1
+
+        else:
+            sentence_features_combined = None
+
+            # Without sentence-level features, the feature sequence lengths are
+            # completely determined by sequence-level features.
+            combined_sequence_sentence_feature_lengths = sequence_feature_lengths
+
+        return sentence_features_combined, combined_sequence_sentence_feature_lengths
+
+    def call(
+        self,
+        inputs: Tuple[
+            List[Union[tf.Tensor, tf.SparseTensor]],
+            List[Union[tf.Tensor, tf.SparseTensor]],
+            tf.Tensor,
+        ],
+        training: bool = False,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Combines multiple 3-D dense/sparse feature tensors into one.
+
+        Arguments:
+            inputs: Tuple containing:
+                sequence_features: Dense or sparse tensors representing different
+                    token-level features.
+                sentence_features: Dense or sparse tensors representing sentence-level
+                    features.
+                sequence_feature_lengths: A tensor containing the real sequence length
+                    (the number of real -- not padding -- tokens) for each example in
+                    the batch.
+            training: A flag indicating whether the layer should behave in training mode
+                (applying dropout to sparse tensors if applicable) or in inference mode
+                (not applying dropout).
+
+        Returns:
+            combined features: A tensor containing all the features combined.
+            mask_combined_sequence_sentence: A binary mask with 1s in place of real
+                features in the combined feature tensor, and 0s in padded positions with
+                fake features.
+        """
+        sequence_features = inputs[0]
+        sentence_features = inputs[1]
+        sequence_feature_lengths = inputs[2]
+
+        # This mask is specifically for sequence-level features.
+        mask_sequence = compute_mask(sequence_feature_lengths)
+
+        sequence_features_combined = self._combine_sequence_level_features(
+            sequence_features, mask_sequence, training
+        )
+
+        (
+            sentence_features_combined,
+            combined_sequence_sentence_feature_lengths,
+        ) = self._combine_sentence_level_features(
+            sentence_features, sequence_feature_lengths, training
+        )
+
+        mask_combined_sequence_sentence = compute_mask(
+            combined_sequence_sentence_feature_lengths
+        )
+
+        # If both feature types are present, combine them. Otherwise, just the present
+        # feature type will be returned.
+        if (
+            sequence_features_combined is not None
+            and sentence_features_combined is not None
+        ):
+            features_to_return = self._concat_sequence_sentence_features(
+                sequence_features_combined,
+                sentence_features_combined,
+                mask_combined_sequence_sentence,
+            )
+        elif sequence_features_combined is not None:
+            features_to_return = sequence_features_combined
+        else:
+            features_to_return = sentence_features_combined
+
+        return features_to_return, mask_combined_sequence_sentence
+
+
+class RasaSequenceLayer(RasaCustomLayer):
+    """Creates an embedding from all features for a sequence attribute; facilitates MLM.
+
+    This layer combines all features for an attribute and embeds them using a
+    transformer, optionally doing masked language modeling. The layer is meant only for
+    attributes with sequence-level features, such as `text`, `response` and
+    `action_text`.
+
+    Internally, this layer applies the following steps:
+    1. Combine features using `RasaFeatureCombiningLayer`.
+    2. Apply a dense layer(s) to the combined features.
+    3. Optionally, and only during training for the `text` attribute, apply masking to
+        the features and create further helper variables for masked language modeling.
+    4. Embed the features using a transformer, effectively reducing variable-length
+        sequences of features to fixed-size embeddings.
+
+    Arguments:
+        attribute: Name of attribute (e.g. `text` or `label`) whose features will be
+            processed.
+        attribute_signature: A dictionary containing two lists of feature signatures,
+            one for each feature type (`sentence` or `sequence`) of the given attribute.
+        config: A model config used for correctly parameterising the underlying layers.
+
+    Input shape:
+        Tuple of three input tensors:
+            sequence_features: List of 3-D dense or sparse tensors, each with shape
+                `(batch_size, max_seq_length, input_dim)` where `input_dim` can be
+                different for sparse vs dense tensors. See the input shape of
+                `ConcatenateSparseDenseFeatures` for more information.
+            sentence_features: List of 3-D dense or sparse tensors, each with shape
+                `(batch_size, 1, input_dim)` where `input_dim` can be different for
+                sparse vs dense tensors, and can differ from that in
+                `sequence_features`. See the input shape of
+                `ConcatenateSparseDenseFeatures` for more information.
+            sequence_feature_lengths: Dense tensor of shape `(batch_size, )`.
+
+    Output shape:
+        outputs: `(batch_size, seq_length, units)` where `units` matches the underlying
+            transformer's output size (if present), otherwise it matches the output size
+            of the `Ffnn` block applied to the combined features, or it's the output
+            size of the underlying `RasaFeatureCombiningLayer` if the `Ffnn` block has 0
+            layers. `seq_length` is the sum of the sequence dimension
+            sizes of sequence- and sentence-level features (for details, see the output
+            shape of `RasaFeatureCombiningLayer`). If both feature types are present,
+            then `seq_length` will be 1 + the length of the longest sequence of real
+            tokens across all examples in the given batch.
+        seq_sent_features: `(batch_size, seq_length, hidden_dim)`, where `hidden_dim` is
+            the output size of the underlying `Ffnn` block, or the output size of the
+            underlying `RasaFeatureCombiningLayer` if the `Ffnn` block has 0 layers.
+        mask_combined_sequence_sentence: `(batch_size, seq_length, 1)`
+        token_ids: `(batch_size, seq_length, id_dim)`. `id_dim` is 2 when no dense
+            sequence-level features are present. Otherwise, it's arbitrarily chosen to
+            match the last dimension size of the first dense sequence-level feature in
+            the input list of features.
+        mlm_boolean_mask: `(batch_size, seq_length, 1)`, empty tensor if not doing MLM.
+        attention_weights: `(transformer_layers, batch_size, num_transformer_heads,
+            seq_length, seq_length)`, empty tensor if the transformer has 0 layers.
+
+    Raises:
+        A `TFLayerConfigException` if no feature signatures for sequence-level features
+            are provided.
+
+    Attributes:
+        output_units: The last dimension size of the layer's first output (`outputs`).
+    """
+
+    FEATURE_COMBINING = "feature_combining"
+    FFNN = "ffnn"
+    TRANSFORMER = "transformer"
+    MLM_INPUT_MASK = "mlm_input_mask"
+    SPARSE_TO_DENSE_FOR_TOKEN_IDS = "sparse_to_dense_for_token_ids"
+
+    def __init__(
+        self,
+        attribute: Text,
+        attribute_signature: Dict[Text, List[FeatureSignature]],
+        config: Dict[Text, Any],
+    ) -> None:
+        """Creates a new `RasaSequenceLayer` object."""
+        if not attribute_signature or not attribute_signature.get(SEQUENCE, []):
+            raise TFLayerConfigException(
+                "The attribute signature must contain some sequence-level feature"
+                "signatures but none were found."
+            )
+
+        super().__init__(name=f"rasa_sequence_layer_{attribute}")
+
+        self._tf_layers: Dict[Text, Any] = {
+            self.FEATURE_COMBINING: RasaFeatureCombiningLayer(
+                attribute, attribute_signature, config
+            ),
+            self.FFNN: layers.Ffnn(
+                config[HIDDEN_LAYERS_SIZES][attribute],
+                config[DROP_RATE],
+                config[REGULARIZATION_CONSTANT],
+                config[CONNECTION_DENSITY],
+                layer_name_suffix=attribute,
+            ),
+        }
+
+        self._enables_mlm = False
+        # Note: Within TED, masked language modeling becomes just input dropout,
+        # since there is no loss term associated with predicting the masked tokens.
+        self._prepare_masked_language_modeling(attribute, attribute_signature, config)
+
+        transformer_layers, transformer_units = self._prepare_transformer(
+            attribute, config
+        )
+        self._has_transformer = transformer_layers > 0
+
+        self.output_units = self._calculate_output_units(
+            attribute, transformer_layers, transformer_units, config
+        )
+
+    @staticmethod
+    def _get_transformer_dimensions(
+        attribute: Text, config: Dict[Text, Any]
+    ) -> Tuple[int, int]:
+        """Determines # of transformer layers & output size from the model config.
+
+        The config can contain these directly (same for all attributes) or specified
+        separately for each attribute.
+        If a transformer is used (e.i. if `number_of_transformer_layers` is positive),
+        the default `transformer_size` which is `None` breaks things. Thus,
+        we need to set a reasonable default value so that the model works fine.
+        """
+        transformer_layers = config[NUM_TRANSFORMER_LAYERS]
+        if isinstance(transformer_layers, dict):
+            transformer_layers = transformer_layers[attribute]
+        transformer_units = config[TRANSFORMER_SIZE]
+        if isinstance(transformer_units, dict):
+            transformer_units = transformer_units[attribute]
+        if transformer_layers > 0 and (not transformer_units or transformer_units < 1):
+            transformer_units = DEFAULT_TRANSFORMER_SIZE
+
+        return transformer_layers, transformer_units
+
+    def _prepare_transformer(
+        self, attribute: Text, config: Dict[Text, Any]
+    ) -> Tuple[int, int]:
+        """Creates a transformer & returns its number of layers and output units."""
+        transformer_layers, transformer_units = self._get_transformer_dimensions(
+            attribute, config
+        )
+        self._tf_layers[self.TRANSFORMER] = prepare_transformer_layer(
+            attribute_name=attribute,
+            config=config,
+            num_layers=transformer_layers,
+            units=transformer_units,
+            drop_rate=config[DROP_RATE],
+            unidirectional=config[UNIDIRECTIONAL_ENCODER],
+        )
+        return transformer_layers, transformer_units
+
+    def _prepare_masked_language_modeling(
+        self,
+        attribute: Text,
+        attribute_signature: Dict[Text, List[FeatureSignature]],
+        config: Dict[Text, Any],
+    ) -> None:
+        """Prepares masking and computing helper variables for masked language modeling.
+
+        Only done for the text attribute and only if sequence-level (token-level)
+        features are present (MLM requires token-level information).
+        """
+        if attribute == TEXT and SEQUENCE in attribute_signature and config[MASKED_LM]:
+            self._enables_mlm = True
+            self._tf_layers[self.MLM_INPUT_MASK] = layers.InputMask()
+
+            # Unique IDs of different token types are needed to construct the possible
+            # label space for MLM. If dense features are present, they're used as such
+            # IDs, othwerise sparse features are embedded by a non-trainable
+            # DenseForSparse layer to create small embeddings that serve as IDs.
+            expect_dense_seq_features = any(
+                [not signature.is_sparse for signature in attribute_signature[SEQUENCE]]
+            )
+            if not expect_dense_seq_features:
+                self._tf_layers[self.SPARSE_TO_DENSE_FOR_TOKEN_IDS] = (
+                    layers.DenseForSparse(
+                        units=2,
+                        use_bias=False,
+                        trainable=False,
+                        name=f"{self.SPARSE_TO_DENSE_FOR_TOKEN_IDS}.{attribute}",
+                    )
+                )
+
+    def _calculate_output_units(
+        self,
+        attribute: Text,
+        transformer_layers: int,
+        transformer_units: int,
+        config: Dict[Text, Any],
+    ) -> int:
+        """Determines the output units based on what layer components are present.
+
+        The size depends on which component is the last created one in the internal
+        pipeline that is `RasaFeatureCombiningLayer` -> `Ffnn` -> `Transformer`, since
+        not all the components are always created.
+        """
+        # transformer is the last component
+        if transformer_layers > 0:
+            return transformer_units
+
+        # the Ffnn block is the last component
+        if len(config[HIDDEN_LAYERS_SIZES][attribute]) > 0:
+            # this is the output size of the last layer of the Ffnn block
+            return config[HIDDEN_LAYERS_SIZES][attribute][-1]
+
+        # only the RasaFeatureCombiningLayer is present
+        return self._tf_layers[self.FEATURE_COMBINING].output_units
+
+    def _features_as_token_ids(
+        self, features: List[Union[tf.Tensor, tf.SparseTensor]]
+    ) -> Optional[tf.Tensor]:
+        """Creates dense labels (token IDs) used for negative sampling in MLM."""
+        # If there are dense features, we use them as labels - taking the first dense
+        # feature in the list, but any other dense feature would do the job.
+        for f in features:
+            if not isinstance(f, tf.SparseTensor):
+                return tf.stop_gradient(f)
+
+        # If no dense features are found, use a sparse feature but convert it into
+        # a dense one first.
+        for f in features:
+            if isinstance(f, tf.SparseTensor):
+                return tf.stop_gradient(
+                    self._tf_layers[self.SPARSE_TO_DENSE_FOR_TOKEN_IDS](f)
+                )
+
+        return None
+
+    def _create_mlm_tensors(
+        self,
+        sequence_features: List[Union[tf.Tensor, tf.SparseTensor]],
+        seq_sent_features: tf.Tensor,
+        mask_sequence: tf.Tensor,
+        sentence_features_present: bool,
+        training: bool,
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        """Produces helper variables for masked language modelling (only in training).
+
+        The `token_ids` embeddings can be viewed as token-level labels/unique IDs of all
+        input tokens (to be used later in the MLM loss) because these embeddings aren't
+        affected by dropout or masking and are effectively always unique for different
+        input tokens (and same for the same tokens).
+        `token_ids` share the batch and sequence dimension with the combined sequence-
+        and sentence-level features, the last dimension is unimportant and mimics the
+        first dense sequence-level feature in the list of features, or alternatively the
+        last dimension will have size 2 if there are only sparse sequence features
+        present.
+        """
+        token_ids = self._features_as_token_ids(sequence_features)
+
+        # Pad in the sequence dimension to match the shape of combined sequence- and
+        # sentence-level features. This means padding by 1 if sentence-level features
+        # are present (those effectively have sequence length of 1) and not padding
+        # otherwise.
+        if sentence_features_present:
+            token_ids = tf.pad(token_ids, [[0, 0], [0, 1], [0, 0]])
+            mask_sequence = tf.pad(mask_sequence, [[0, 0], [0, 1], [0, 0]])
+
+        # mlm_boolean_mask has the same shape as the tensor with all combined features
+        # (except the last dimension), with True meaning tokens that are masked and
+        # False meaning tokens that aren't masked or that are fake (padded) tokens.
+        # Note that only sequence-level features are masked, nothing happens to the
+        # sentence-level features in the combined features tensor.
+        seq_sent_features, mlm_boolean_mask = self._tf_layers[self.MLM_INPUT_MASK](
+            seq_sent_features, mask_sequence, training
+        )
+
+        return seq_sent_features, token_ids, mlm_boolean_mask
+
+    def call(
+        self,
+        inputs: Tuple[
+            List[Union[tf.Tensor, tf.SparseTensor]],
+            List[Union[tf.Tensor, tf.SparseTensor]],
+            tf.Tensor,
+        ],
+        training: bool = False,
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
+        """Combines all of an attribute's features and embeds using a transformer.
+
+        Arguments:
+            inputs: Tuple containing:
+                sequence_features: Dense or sparse tensors representing different
+                    token-level features.
+                sentence_features: Dense or sparse tensors representing different
+                    sentence-level features.
+                sequence_feature_lengths: A tensor containing the real sequence length
+                    (the number of real -- not padding -- tokens) for each example in
+                    the batch.
+            training: A flag indicating whether the layer should behave in training mode
+                (applying dropout to sparse tensors if applicable) or in inference mode
+                (not applying dropout).
+
+        Returns:
+            outputs: Tensor with all features combined, masked (if doing MLM) and
+                embedded with a transformer.
+            seq_sent_features: Tensor with all features combined from just before the
+                masking and transformer is applied
+            mask_combined_sequence_sentence: A binary mask with 1s in place of real
+                features in the combined feature tensor, and 0s in padded positions with
+                fake features.
+            token_ids: Tensor with dense token-level features which can serve as
+                IDs (unique embeddings) of all the different tokens found in the batch.
+                Empty tensor if not doing MLM.
+            mlm_boolean_mask: A boolean mask with `True` where real tokens in `outputs`
+                were masked and `False` elsewhere. Empty tensor if not doing MLM.
+            attention_weights: Tensor containing self-attention weights received
+                from the underlying transformer. Empty tensor if the transformer has 0
+                layers.
+        """
+        sequence_features = inputs[0]
+        sentence_features = inputs[1]
+        sequence_feature_lengths = inputs[2]
+
+        # Combine all features (sparse/dense, sequence-/sentence-level) into one tensor,
+        # also get a binary mask that has 1s at positions with real features and 0s at
+        # padded positions.
+        seq_sent_features, mask_combined_sequence_sentence = self._tf_layers[
+            self.FEATURE_COMBINING
+        ]((sequence_features, sentence_features, sequence_feature_lengths))
+
+        # Apply one or more dense layers.
+        seq_sent_features = self._tf_layers[self.FFNN](seq_sent_features, training)
+
+        # If using masked language modeling, mask the transformer inputs and get labels
+        # for the masked tokens and a boolean mask. Note that TED does not use MLM loss,
+        # hence using masked language modeling (if enabled) becomes just input dropout.
+        if self._enables_mlm and training:
+            mask_sequence = compute_mask(sequence_feature_lengths)
+            (
+                seq_sent_features_masked,
+                token_ids,
+                mlm_boolean_mask,
+            ) = self._create_mlm_tensors(
+                sequence_features,
+                seq_sent_features,
+                mask_sequence,
+                sentence_features_present=len(sentence_features) > 0,
+                training=training,
+            )
+        else:
+            # tf.zeros((0,)) is an alternative to None
+            token_ids = tf.zeros((0,))
+            mlm_boolean_mask = tf.zeros((0,))
+            seq_sent_features_masked = seq_sent_features
+
+        # Apply the transformer (if present), hence reducing a sequences of features per
+        # input example into a simple fixed-size embedding.
+        if self._has_transformer:
+            mask_padding = 1 - mask_combined_sequence_sentence
+            outputs, attention_weights = self._tf_layers[self.TRANSFORMER](
+                seq_sent_features_masked, mask_padding, training
+            )
+            outputs = tf.nn.gelu(outputs)
+        else:
+            # tf.zeros((0,)) is an alternative to None
+            outputs, attention_weights = seq_sent_features_masked, tf.zeros((0,))
+
+        return (
+            outputs,
+            seq_sent_features,
+            mask_combined_sequence_sentence,
+            token_ids,
+            mlm_boolean_mask,
+            attention_weights,
+        )
+
+
+def compute_mask(sequence_lengths: tf.Tensor) -> tf.Tensor:
+    """Computes binary mask given real sequence lengths.
+
+    Takes a 1-D tensor of shape `(batch_size,)` containing the lengths of sequences
+    (in terms of number of tokens) in the batch. Creates a binary mask of shape
+    `(batch_size, max_seq_length, 1)` with 1s at positions with real tokens and 0s
+    elsewhere.
+    """
+    mask = tf.sequence_mask(sequence_lengths, dtype=tf.float32)
+    return tf.expand_dims(mask, -1)
+
+
+def prepare_transformer_layer(
+    attribute_name: Text,
+    config: Dict[Text, Any],
+    num_layers: int,
+    units: int,
+    drop_rate: float,
+    unidirectional: bool,
+) -> Union[
+    TransformerEncoder,
+    Callable[
+        [tf.Tensor, Optional[tf.Tensor], Optional[Union[tf.Tensor, bool]]],
+        Tuple[tf.Tensor, Optional[tf.Tensor]],
+    ],
+]:
+    """Creates & returns a transformer encoder, potentially with 0 layers."""
+    if num_layers > 0:
+        return TransformerEncoder(
+            num_layers,
+            units,
+            config[NUM_HEADS],
+            units * 4,
+            config[REGULARIZATION_CONSTANT],
+            dropout_rate=drop_rate,
+            attention_dropout_rate=config[DROP_RATE_ATTENTION],
+            density=config[CONNECTION_DENSITY],
+            unidirectional=unidirectional,
+            use_key_relative_position=config[KEY_RELATIVE_ATTENTION],
+            use_value_relative_position=config[VALUE_RELATIVE_ATTENTION],
+            max_relative_position=config[MAX_RELATIVE_POSITION],
+            name=f"{attribute_name}_encoder",
+        )
+    # create lambda so that it can be used later without the check
+    return lambda x, mask, training: (x, None)