| 15 | |
| 16 | |
| 17 | class GGMLTensor: |
| 18 | def __init__( |
| 19 | self, |
| 20 | data: Union[torch.Tensor, np.ndarray, None] = None, |
| 21 | orig_shape: Tuple[int, ...] = None, |
| 22 | dtype: torch.dtype = None, |
| 23 | gguf_type: gguf.GGMLQuantizationType = None, |
| 24 | requires_grad: bool = False, |
| 25 | aligned: bool = True, |
| 26 | pin_memory: bool = False, |
| 27 | preallocated: bool = False, |
| 28 | ): |
| 29 | super().__init__() |
| 30 | |
| 31 | assert orig_shape is not None |
| 32 | assert gguf_type is not None |
| 33 | |
| 34 | if isinstance(data, np.ndarray): |
| 35 | import warnings |
| 36 | |
| 37 | with warnings.catch_warnings(): |
| 38 | warnings.filterwarnings("ignore", message="The given NumPy array is not writable") |
| 39 | torch_data = torch.from_numpy(data) |
| 40 | else: |
| 41 | torch_data = data |
| 42 | |
| 43 | if dtype is not None and torch_data.dtype != dtype: |
| 44 | torch_data = torch_data.to(dtype) |
| 45 | |
| 46 | self.data = torch_data |
| 47 | |
| 48 | self.gguf_type = gguf_type |
| 49 | self._orig_shape = orig_shape |
| 50 | self._aligned = aligned |
| 51 | self._pinned_memory = pin_memory |
| 52 | self._requires_grad = requires_grad |
| 53 | self._preallocated = preallocated |
| 54 | |
| 55 | self._quantized = self._is_quantized_type(gguf_type) |
| 56 | self._q_type = self._get_quant_type_str(gguf_type) |
| 57 | |
| 58 | if aligned: |
| 59 | self._make_aligned() |
| 60 | if pin_memory: |
| 61 | self._pin_memory() |
| 62 | |
| 63 | def _is_quantized_type(self, gguf_type: gguf.GGMLQuantizationType) -> bool: |
| 64 | return gguf_type not in TORCH_COMPATIBLE_QTYPES |
| 65 | |
| 66 | def _get_quant_type_str(self, gguf_type: gguf.GGMLQuantizationType) -> str: |
| 67 | type_mapping = { |
| 68 | gguf.GGMLQuantizationType.F32: "ggml_f32", |
| 69 | gguf.GGMLQuantizationType.F16: "ggml_f16", |
| 70 | gguf.GGMLQuantizationType.Q4_0: "ggml_q4_0", |
| 71 | gguf.GGMLQuantizationType.Q4_1: "ggml_q4_1", |
| 72 | gguf.GGMLQuantizationType.Q5_0: "ggml_q5_0", |
| 73 | gguf.GGMLQuantizationType.Q5_1: "ggml_q5_1", |
| 74 | gguf.GGMLQuantizationType.Q8_0: "ggml_q8_0", |