RasaModelData can handle 4D Tensors

rasa/utils/tensorflow/model_data.py

@@ -63,6 +63,11 @@ def __new__(
 
         return feature_array
 
+    def __init__(self, input_array: Any, number_of_dimensions: int, **kwargs):
+        # Needed in order to avoid 'Invalid keyword argument number_of_dimensions to function FeatureArray.__init__ '
+        super().__init__(**kwargs)
+        self.number_of_dimensions = number_of_dimensions
+
     def __array_finalize__(self, obj: Any) -> None:
         if obj is None:
             return
@@ -152,7 +157,7 @@ def _validate_number_of_dimensions(
                 f"array: {input_array}."
             )
 
-    def get_shape_type_info(self) -> Tuple[List[Tuple], List[Tuple]]:
+    def get_shape_type_info(self) -> Tuple[List, List]:
         """Returns the shape and type information needed to convert this feature array into tensors.
 
         Returns:
@@ -183,7 +188,7 @@ def get_shape_type_info(self) -> Tuple[List[Tuple], List[Tuple]]:
             return [(None, None, None, self.units)], [tf.float32]
 
         return [], []
-        
+
 
 class FeatureSignature(NamedTuple):
     """Stores the shape, the type (sparse vs dense), and the number of dimensions of features."""

rasa/utils/tensorflow/model_data_utils.py

@@ -96,7 +96,7 @@ def convert_to_data_format(
         List[List[Dict[Text, List["Features"]]]], List[Dict[Text, List["Features"]]]
     ],
     zero_state_features: Optional[Dict[Text, List["Features"]]] = None,
-) -> Tuple[Data, Optional[Dict[Text, List[FeatureArray]]]]:
+) -> Tuple[Data, Optional[Dict[Text, List["Features"]]]]:
     """Converts the input into "Data" format.
 
     Args:
@@ -186,8 +186,8 @@ def _features_for_attribute(
         tracker_features, zero_state_features[attribute]
     )
 
-    sparse_features = defaultdict(list)
-    dense_features = defaultdict(list)
+    sparse_features = defaultdict()
+    dense_features = defaultdict()
 
     # vstack serves as removing dimension in case we are not dealing with a sequence
     for key, values in _sparse_features.items():