RasaModelData can handle 4D Tensors

rasa/core/policies/ensemble.py

@@ -29,9 +29,7 @@
     ACTION_BACK_NAME,
 )
 from rasa.shared.core.domain import InvalidDomain, Domain
-from rasa.shared.core.events import (
-    DefinePrevUserUtteredFeaturization,
-)
+from rasa.shared.core.events import DefinePrevUserUtteredFeaturization
 from rasa.shared.core.events import ActionExecutionRejected
 from rasa.core.exceptions import UnsupportedDialogueModelError
 from rasa.core.featurizers.tracker_featurizers import MaxHistoryTrackerFeaturizer

rasa/core/policies/ted_policy.py

@@ -26,7 +26,11 @@
 from rasa.shared.core.generator import TrackerWithCachedStates
 from rasa.utils import train_utils
 from rasa.utils.tensorflow.models import RasaModel, TransformerRasaModel
-from rasa.utils.tensorflow.model_data import RasaModelData, FeatureSignature
+from rasa.utils.tensorflow.model_data import (
+    RasaModelData,
+    FeatureSignature,
+    FeatureArray,
+)
 from rasa.utils.tensorflow.model_data_utils import convert_to_data_format
 from rasa.utils.tensorflow.constants import (
     LABEL,
@@ -264,7 +268,9 @@ def _create_label_data(
 
         label_ids = np.arange(domain.num_actions)
         label_data.add_features(
-            LABEL_KEY, LABEL_SUB_KEY, [np.expand_dims(label_ids, -1)]
+            LABEL_KEY,
+            LABEL_SUB_KEY,
+            [FeatureArray(np.expand_dims(label_ids, -1), number_of_dimensions=2)],
         )
 
         return label_data, encoded_all_labels
@@ -295,7 +301,11 @@ def _create_model_data(
             label_ids = np.array(
                 [np.expand_dims(seq_label_ids, -1) for seq_label_ids in label_ids]
             )
-            model_data.add_features(LABEL_KEY, LABEL_SUB_KEY, [label_ids])
+            model_data.add_features(
+                LABEL_KEY,
+                LABEL_SUB_KEY,
+                [FeatureArray(label_ids, number_of_dimensions=3)],
+            )
 
             attribute_data, self.zero_state_features = convert_to_data_format(
                 tracker_state_features

rasa/nlu/classifiers/diet_classifier.py

@@ -24,7 +24,11 @@
 from rasa.utils import train_utils
 from rasa.utils.tensorflow import layers
 from rasa.utils.tensorflow.models import RasaModel, TransformerRasaModel
-from rasa.utils.tensorflow.model_data import RasaModelData, FeatureSignature
+from rasa.utils.tensorflow.model_data import (
+    RasaModelData,
+    FeatureSignature,
+    FeatureArray,
+)
 from rasa.nlu.constants import TOKENS_NAMES
 from rasa.shared.nlu.constants import (
     TEXT,
@@ -325,7 +329,7 @@ def __init__(
         self.model = model
 
         self._label_data: Optional[RasaModelData] = None
-        self._data_example: Optional[Dict[Text, List[np.ndarray]]] = None
+        self._data_example: Optional[Dict[Text, List[FeatureArray]]] = None
 
     @property
     def label_key(self) -> Optional[Text]:
@@ -492,10 +496,10 @@ def _check_input_dimension_consistency(self, model_data: RasaModelData) -> None:
         """Checks if features have same dimensionality if hidden layers are shared."""
 
         if self.component_config.get(SHARE_HIDDEN_LAYERS):
-            num_text_sentence_features = model_data.feature_dimension(TEXT, SENTENCE)
-            num_label_sentence_features = model_data.feature_dimension(LABEL, SENTENCE)
-            num_text_sequence_features = model_data.feature_dimension(TEXT, SEQUENCE)
-            num_label_sequence_features = model_data.feature_dimension(LABEL, SEQUENCE)
+            num_text_sentence_features = model_data.number_of_units(TEXT, SENTENCE)
+            num_label_sentence_features = model_data.number_of_units(LABEL, SENTENCE)
+            num_text_sequence_features = model_data.number_of_units(TEXT, SEQUENCE)
+            num_label_sequence_features = model_data.number_of_units(LABEL, SEQUENCE)
 
             if (0 < num_text_sentence_features != num_label_sentence_features > 0) or (
                 0 < num_text_sequence_features != num_label_sequence_features > 0
@@ -507,7 +511,7 @@ def _check_input_dimension_consistency(self, model_data: RasaModelData) -> None:
 
     def _extract_labels_precomputed_features(
         self, label_examples: List[Message], attribute: Text = INTENT
-    ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
+    ) -> Tuple[List[FeatureArray], List[FeatureArray]]:
         """Collects precomputed encodings."""
 
         features = defaultdict(list)
@@ -521,23 +525,32 @@ def _extract_labels_precomputed_features(
         sentence_features = []
         for feature_name, feature_value in features.items():
             if SEQUENCE in feature_name:
-                sequence_features.append(np.array(features[feature_name]))
+                sequence_features.append(
+                    FeatureArray(np.array(feature_value), number_of_dimensions=3)
+                )
             else:
-                sentence_features.append(np.array(features[feature_name]))
+                sentence_features.append(
+                    FeatureArray(np.array(feature_value), number_of_dimensions=3)
+                )
 
-        return (sequence_features, sentence_features)
+        return sequence_features, sentence_features
 
     @staticmethod
     def _compute_default_label_features(
         labels_example: List[Message],
-    ) -> List[np.ndarray]:
+    ) -> List[FeatureArray]:
         """Computes one-hot representation for the labels."""
 
         logger.debug("No label features found. Computing default label features.")
 
         eye_matrix = np.eye(len(labels_example), dtype=np.float32)
         # add sequence dimension to one-hot labels
-        return [np.array([np.expand_dims(a, 0) for a in eye_matrix])]
+        return [
+            FeatureArray(
+                np.array([np.expand_dims(a, 0) for a in eye_matrix]),
+                number_of_dimensions=3,
+            )
+        ]
 
     def _create_label_data(
         self,
@@ -590,16 +603,23 @@ def _create_label_data(
         # explicitly add last dimension to label_ids
         # to track correctly dynamic sequences
         label_data.add_features(
-            LABEL_KEY, LABEL_SUB_KEY, [np.expand_dims(label_ids, -1)]
+            LABEL_KEY,
+            LABEL_SUB_KEY,
+            [FeatureArray(np.expand_dims(label_ids, -1), number_of_dimensions=2)],
         )
 
         label_data.add_lengths(LABEL, SEQUENCE_LENGTH, LABEL, SEQUENCE)
 
         return label_data
 
-    def _use_default_label_features(self, label_ids: np.ndarray) -> List[np.ndarray]:
+    def _use_default_label_features(self, label_ids: np.ndarray) -> List[FeatureArray]:
         all_label_features = self._label_data.get(LABEL, SENTENCE)[0]
-        return [np.array([all_label_features[label_id] for label_id in label_ids])]
+        return [
+            FeatureArray(
+                np.array([all_label_features[label_id] for label_id in label_ids]),
+                number_of_dimensions=2,
+            )
+        ]
 
     def _create_model_data(
         self,
@@ -645,7 +665,11 @@ def _create_model_data(
         for key, attribute_features in features.items():
             for sub_key, _features in attribute_features.items():
                 sub_key = sub_key.replace(f"{SPARSE}_", "").replace(f"{DENSE}_", "")
-                model_data.add_features(key, sub_key, [np.array(_features)])
+                model_data.add_features(
+                    key,
+                    sub_key,
+                    [FeatureArray(np.array(_features), number_of_dimensions=3)],
+                )
 
         if (
             label_attribute
@@ -660,7 +684,9 @@ def _create_model_data(
         # explicitly add last dimension to label_ids
         # to track correctly dynamic sequences
         model_data.add_features(
-            LABEL_KEY, LABEL_SUB_KEY, [np.expand_dims(label_ids, -1)]
+            LABEL_KEY,
+            LABEL_SUB_KEY,
+            [FeatureArray(np.expand_dims(label_ids, -1), number_of_dimensions=2)],
         )
 
         model_data.add_lengths(TEXT, SEQUENCE_LENGTH, TEXT, SEQUENCE)
@@ -1028,7 +1054,7 @@ def _load_model(
         entity_tag_specs: List[EntityTagSpec],
         label_data: RasaModelData,
         meta: Dict[Text, Any],
-        data_example: Dict[Text, Dict[Text, List[np.ndarray]]],
+        data_example: Dict[Text, Dict[Text, List[FeatureArray]]],
         model_dir: Text,
     ) -> "RasaModel":
         file_name = meta.get("file")