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}." )