chore: import upstream snapshot with attribution

rasa/utils/tensorflow/feature_array.py

@@ -0,0 +1,370 @@
+from typing import Dict, Any, List, Tuple, Optional, Union
+
+import numpy as np
+import scipy.sparse
+from safetensors.numpy import load_file
+from safetensors.numpy import save_file
+
+import rasa.shared.utils.io
+
+
+def _recursive_serialize(
+    array: Any, prefix: str, data_dict: Dict[str, Any], metadata: List[Dict[str, Any]]
+) -> None:
+    """Recursively serialize arrays and matrices for high dimensional data."""
+    if isinstance(array, np.ndarray) and array.ndim <= 2:
+        data_key = f"{prefix}_array"
+        data_dict[data_key] = array
+        metadata.append({"type": "dense", "key": data_key, "shape": array.shape})
+
+    elif isinstance(array, list) and all([isinstance(v, float) for v in array]):
+        data_key = f"{prefix}_list"
+        data_dict[data_key] = np.array(array, dtype=np.float32)
+        metadata.append({"type": "list", "key": data_key})
+
+    elif isinstance(array, list) and all([isinstance(v, int) for v in array]):
+        data_key = f"{prefix}_list"
+        data_dict[data_key] = np.array(array, dtype=np.int64)
+        metadata.append({"type": "list", "key": data_key})
+
+    elif isinstance(array, scipy.sparse.spmatrix):
+        data_key_data = f"{prefix}_data"
+        data_key_row = f"{prefix}_row"
+        data_key_col = f"{prefix}_col"
+        array = array.tocoo()
+        data_dict.update(
+            {
+                data_key_data: array.data,
+                data_key_row: array.row,
+                data_key_col: array.col,
+            }
+        )
+        metadata.append({"type": "sparse", "key": prefix, "shape": array.shape})
+
+    elif isinstance(array, list) or isinstance(array, np.ndarray):
+        group_metadata = {"type": "group", "subcomponents": []}
+        for idx, item in enumerate(array):
+            new_prefix = f"{prefix}_{idx}"
+            _recursive_serialize(
+                item, new_prefix, data_dict, group_metadata["subcomponents"]
+            )
+        metadata.append(group_metadata)
+
+
+def _serialize_nested_data(
+    nested_data: Dict[str, Dict[str, List["FeatureArray"]]],
+    prefix: str,
+    data_dict: Dict[str, np.ndarray],
+    metadata: List[Dict[str, Union[str, List]]],
+) -> None:
+    """Handle serialization across dictionary and list levels."""
+    for outer_key, inner_dict in nested_data.items():
+        inner_metadata = {"key": outer_key, "components": []}
+
+        for inner_key, feature_arrays in inner_dict.items():
+            array_metadata = {
+                "key": inner_key,
+                "number_of_dimensions": feature_arrays[0].number_of_dimensions,
+                "features": [],
+            }
+
+            for idx, feature_array in enumerate(feature_arrays):
+                feature_prefix = f"{prefix}_{outer_key}_{inner_key}_{idx}"
+                _recursive_serialize(
+                    feature_array.tolist(),
+                    feature_prefix,
+                    data_dict,
+                    array_metadata["features"],
+                )
+
+            inner_metadata["components"].append(  # type:ignore[attr-defined]
+                array_metadata
+            )
+
+        metadata.append(inner_metadata)
+
+
+def serialize_nested_feature_arrays(
+    nested_feature_array: Dict[str, Dict[str, List["FeatureArray"]]],
+    data_filename: str,
+    metadata_filename: str,
+) -> None:
+    data_dict: Dict[str, np.ndarray] = {}
+    metadata: List[Dict[str, Union[str, List]]] = []
+
+    _serialize_nested_data(nested_feature_array, "component", data_dict, metadata)
+
+    # Save serialized data and metadata
+    save_file(data_dict, data_filename)
+    rasa.shared.utils.io.dump_obj_as_json_to_file(metadata_filename, metadata)
+
+
+def _recursive_deserialize(
+    metadata: List[Dict[str, Any]], data: Dict[str, Any]
+) -> List[Any]:
+    """Recursively deserialize arrays and matrices for high dimensional data."""
+    result = []
+
+    for item in metadata:
+        if item["type"] == "dense":
+            key = item["key"]
+            array = np.asarray(data[key]).reshape(item["shape"])
+            result.append(array)
+
+        elif item["type"] == "list":
+            key = item["key"]
+            result.append(list(data[key]))
+
+        elif item["type"] == "sparse":
+            data_vals = data[f"{item['key']}_data"]
+            row_vals = data[f"{item['key']}_row"]
+            col_vals = data[f"{item['key']}_col"]
+            sparse_matrix = scipy.sparse.coo_matrix(
+                (data_vals, (row_vals, col_vals)), shape=item["shape"]
+            )
+            result.append(sparse_matrix)
+        elif item["type"] == "group":
+            sublist = _recursive_deserialize(item["subcomponents"], data)
+            result.append(sublist)
+
+    return result
+
+
+def _deserialize_nested_data(
+    metadata: List[Dict[str, Any]], data_dict: Dict[str, Any]
+) -> Dict[str, Dict[str, List["FeatureArray"]]]:
+    """Handle deserialization across all dictionary and list levels."""
+    result: Dict[str, Dict[str, List["FeatureArray"]]] = {}
+
+    for outer_item in metadata:
+        outer_key = outer_item["key"]
+        result[outer_key] = {}
+
+        for inner_item in outer_item["components"]:
+            inner_key = inner_item["key"]
+            feature_arrays = []
+
+            # Reconstruct the list of FeatureArrays
+            for feature_item in inner_item["features"]:
+                # Reconstruct the list of FeatureArrays
+                feature_array_data = _recursive_deserialize([feature_item], data_dict)
+                # Prepare the input for the FeatureArray;
+                # ensure it is np.ndarray compatible
+                input_array = np.array(feature_array_data[0], dtype=object)
+                feature_array = FeatureArray(
+                    input_array, inner_item["number_of_dimensions"]
+                )
+                feature_arrays.append(feature_array)
+
+            result[outer_key][inner_key] = feature_arrays
+
+    return result
+
+
+def deserialize_nested_feature_arrays(
+    data_filename: str, metadata_filename: str
+) -> Dict[str, Dict[str, List["FeatureArray"]]]:
+    metadata = rasa.shared.utils.io.read_json_file(metadata_filename)
+    data_dict = load_file(data_filename)
+
+    return _deserialize_nested_data(metadata, data_dict)
+
+
+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]
+        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: str, *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 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
+        """
+        # when loading the feature arrays from disk, the shape represents
+        # the correct number of dimensions
+        if len(input_array.shape) == number_of_dimensions:
+            return
+
+        _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}."
+            )