Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/llm_datadist_v1/tensor.py: 89%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# ----------------------------------------------------------------------------------------------------------- 

4# Copyright (c) 2025 Huawei Technologies Co., Ltd. 

5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 

6# CANN Open Software License Agreement Version 2.0 (the "License"). 

7# Please refer to the License for details. You may not use this file except in compliance with the License. 

8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 

9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. 

10# See LICENSE in the root of the software repository for the full text of the License. 

11# ----------------------------------------------------------------------------------------------------------- 

12 

13import ctypes 

14from typing import List, Tuple, Union 

15 

16import numpy as np 

17 

18from llm_datadist_v1 import data_type, llm_wrapper 

19from llm_datadist_v1.data_type import _dwrapper_dtype_to_python_dtype 

20from llm_datadist_v1.status import handle_llm_status 

21from llm_datadist_v1.utils import utils 

22 

23 

24class TensorDesc(object): 

25 def __init__(self, dtype: data_type.DataType, shape: Union[List[int], Tuple[int]]): 

26 """ 

27 初始化 

28 Args: 

29 dtype: 数据类型 

30 shape: 数据维度信息 

31 """ 

32 utils.check_isinstance("dtype", dtype, data_type.DataType) 

33 utils.check_isinstance("shape", shape, [list, tuple], int) 

34 self._dtype = dtype 

35 self._shape = list(shape) 

36 

37 @property 

38 def dtype(self): 

39 return self._dtype 

40 

41 @property 

42 def shape(self): 

43 return self._shape 

44 

45 def __str__(self): 

46 return f"TensorDesc(dtype={str(self.dtype)}, shape={str(self.shape)})" 

47 

48 

49class Tensor(object): 

50 def __init__(self, data, tensor_desc: TensorDesc = None): 

51 """ 

52 初始化 

53 Args: 

54 data: 数据 

55 tensor_desc: 描述信息 

56 """ 

57 utils.check_isinstance("data", data, [np.ndarray, Tensor, int]) 

58 utils.check_isinstance("tensor_desc", tensor_desc, TensorDesc) 

59 self._tensor_id = 0 

60 if utils.check_type(data, Tensor): 

61 self._tensor_desc = data._tensor_desc 

62 self._tensor_id = llm_wrapper.clone_tensor(data._tensor_id) 

63 elif utils.check_type(data, int): 

64 self._tensor_desc = tensor_desc 

65 self._tensor_id = data 

66 else: 

67 self._init_by_ndarray(data, tensor_desc) 

68 

69 def __del__(self): 

70 # 保底释放kv cache, 但更推荐主动通过调用kv_cache_manager.deallocate_cache来释放kv cache，而不应该遗留到此处自动释放 

71 if self._tensor_id != 0: 

72 llm_wrapper.destroy_tensor(self._tensor_id) 

73 

74 def __str__(self): 

75 return f"Tensor({self.numpy(True if self._is_inner_dtype_str() else False)},tensor_desc={self._tensor_desc})" 

76 

77 @staticmethod 

78 def from_tensor_tuple(tensor_tuple: Tuple[int, int, List[int]]): 

79 tensor_desc = TensorDesc(_dwrapper_dtype_to_python_dtype[tensor_tuple[1]], tensor_tuple[2]) 

80 return Tensor(tensor_tuple[0], tensor_desc) 

81 

82 def _init_by_ndarray(self, data: np.ndarray, tensor_desc: TensorDesc = None): 

83 if tensor_desc: 

84 if list(data.shape) != tensor_desc.shape: 

85 raise RuntimeError( 

86 f"The shape of data:{data.shape} is not same as tensor_desc shape:{tensor_desc.shape}" 

87 ) 

88 desc_np_dtype = data_type.dtype_to_np_dtype.get(tensor_desc.dtype) 

89 if data.dtype != desc_np_dtype: 

90 raise RuntimeError( 

91 f"The dtype of data:{data.dtype} is not same as tensor_desc dtype:{tensor_desc.dtype}" 

92 ) 

93 else: 

94 if data.dtype not in data_type.valid_np_dtypes and not self._is_origin_dtype_str(data.dtype): 

95 raise RuntimeError( 

96 f"The dtype of data:{data.dtype} is not valid, only support {data_type.valid_np_dtypes}" 

97 ) 

98 if tensor_desc: 

99 self._tensor_desc = tensor_desc 

100 elif self._is_origin_dtype_str(data.dtype): 

101 self._tensor_desc = TensorDesc(data_type.DataType.DT_STRING, list(data.shape)) 

102 else: 

103 self._tensor_desc = TensorDesc(data_type.np_dtype_to_dtype[data.dtype], list(data.shape)) 

104 if self._is_origin_dtype_str(data.dtype): 

105 data = self._convert_raw_str_data(data) 

106 if not data.flags.c_contiguous: 

107 raise RuntimeError("The data is not c_contiguous") 

108 data_ptr = data.ctypes.data_as(ctypes.c_void_p).value 

109 size = data.nbytes 

110 self._tensor_id = llm_wrapper.build_tensor( 

111 data_ptr, 

112 size, 

113 data_type.python_dtype_2_dwrapper_dtype.get(self._tensor_desc.dtype), 

114 list(self._tensor_desc.shape), 

115 ) 

116 

117 def _is_origin_dtype_str(self, dtype): 

118 return np.issubdtype(dtype, np.str_) or np.issubdtype(dtype, np.bytes_) 

119 

120 def _is_inner_dtype_str(self): 

121 return self._tensor_desc is not None and self._tensor_desc.dtype == data_type.DataType.DT_STRING 

122 

123 def _convert_raw_str_data(self, data): 

124 format_data = data.astype(np.bytes_) 

125 end_point = "\0".encode("ascii", errors="ignore") 

126 new_data = np.char.add(format_data, end_point) 

127 return new_data 

128 

129 def numpy(self, copy=False): 

130 """ 

131 获取数据的numpy表示 

132 Args: 

133 copy: 是否复制 

134 

135 Returns: 

136 数据的numpy表示 

137 """ 

138 utils.check_isinstance("copy", copy, bool) 

139 if self._is_inner_dtype_str(): 

140 if not copy: 

141 raise RuntimeError("String tensor only support when param copy is True.") 

142 return np.array(llm_wrapper.get_string_tensor(self._tensor_id)).reshape(self._tensor_desc.shape) 

143 ret, tensor = llm_wrapper.tensor_get_buffer(self._tensor_id) 

144 handle_llm_status(ret, "Tensor.numpy", "Failed to get tensor buffer") 

145 if self._tensor_desc.dtype == data_type.DataType.DT_BF16: 

146 np_array = np.frombuffer(tensor, dtype=np.uint16) 

147 return (np_array.astype(np.uint32) << 16).view(np.float32) 

148 elif self._tensor_desc.dtype == data_type.DataType.DT_FLOAT16: 

149 np_array = np.frombuffer(tensor, dtype=np.uint16) 

150 return np_array.view(np.float16) 

151 if copy: 

152 ret = np.array(tensor, copy=True) 

153 else: 

154 ret = np.asarray(tensor) 

155 return ret