safetensors 0.6.2__cp38-abi3-win_amd64.whl → 0.7.0rc0__cp38-abi3-win_amd64.whl

safetensors/numpy.py

@@ -154,6 +154,7 @@ _TYPES = {
     "I8": np.int8,
     "U8": np.uint8,
     "BOOL": bool,
+    "C64": np.complex64,
 }
 
 

safetensors/paddle.py

@@ -1,10 +1,11 @@
 import os
-from typing import Dict, Optional, Union
+import sys
+from typing import Any, Dict, Optional, Union
 
 import numpy as np
-
 import paddle
-from safetensors import numpy
+
+from safetensors import numpy, deserialize, safe_open, serialize, serialize_file
 
 
 def save(
@@ -34,8 +35,9 @@ def save(
     byte_data = save(tensors)
     ```
     """
-    np_tensors = _paddle2np(tensors)
-    return numpy.save(np_tensors, metadata=metadata)
+    serialized = serialize(_flatten(tensors), metadata=metadata)
+    result = bytes(serialized)
+    return result
 
 
 def save_file(
@@ -69,8 +71,7 @@ def save_file(
     save_file(tensors, "model.safetensors")
     ```
     """
-    np_tensors = _paddle2np(tensors)
-    return numpy.save_file(np_tensors, filename, metadata=metadata)
+    serialize_file(_flatten(tensors), filename, metadata=metadata)
 
 
 def load(data: bytes, device: str = "cpu") -> Dict[str, paddle.Tensor]:
@@ -96,8 +97,12 @@ def load(data: bytes, device: str = "cpu") -> Dict[str, paddle.Tensor]:
     loaded = load(data)
     ```
     """
-    flat = numpy.load(data)
-    return _np2paddle(flat, device)
+    if paddle.__version__ >= "3.2.0":
+        flat = deserialize(data)
+        return _view2paddle(flat, device)
+    else:
+        flat = numpy.load(data)
+        return _np2paddle(flat, device)
 
 
 def load_file(
@@ -125,9 +130,15 @@ def load_file(
     loaded = load_file(file_path)
     ```
     """
-    flat = numpy.load_file(filename)
-    output = _np2paddle(flat, device)
-    return output
+    result = {}
+    if paddle.__version__ >= "3.2.0":
+        with safe_open(filename, framework="paddle", device=device) as f:
+            for k in f.offset_keys():
+                result[k] = f.get_tensor(k)
+    else:
+        flat = numpy.load_file(filename)
+        result = _np2paddle(flat, device)
+    return result
 
 
 def _np2paddle(
@@ -142,3 +153,138 @@ def _paddle2np(paddle_dict: Dict[str, paddle.Tensor]) -> Dict[str, np.array]:
     for k, v in paddle_dict.items():
         paddle_dict[k] = v.detach().cpu().numpy()
     return paddle_dict
+
+
+_SIZE = {
+    paddle.int64: 8,
+    paddle.float32: 4,
+    paddle.int32: 4,
+    paddle.bfloat16: 2,
+    paddle.float16: 2,
+    paddle.int16: 2,
+    paddle.uint8: 1,
+    paddle.int8: 1,
+    paddle.bool: 1,
+    paddle.float64: 8,
+    paddle.float8_e4m3fn: 1,
+    paddle.float8_e5m2: 1,
+    paddle.complex64: 8,
+    # XXX: These are not supported yet in paddle
+    # paddle.uint64: 8,
+    # paddle.uint32: 4,
+    # paddle.uint16: 2,
+    # paddle.float8_e8m0: 1,
+    # paddle.float4_e2m1_x2: 1,
+}
+
+_TYPES = {
+    "F64": paddle.float64,
+    "F32": paddle.float32,
+    "F16": paddle.float16,
+    "BF16": paddle.bfloat16,
+    "I64": paddle.int64,
+    "I32": paddle.int32,
+    "I16": paddle.int16,
+    "I8": paddle.int8,
+    "U8": paddle.uint8,
+    "BOOL": paddle.bool,
+    "F8_E4M3": paddle.float8_e4m3fn,
+    "F8_E5M2": paddle.float8_e5m2,
+}
+
+NPDTYPES = {
+    paddle.int64: np.int64,
+    paddle.float32: np.float32,
+    paddle.int32: np.int32,
+    # XXX: This is ok because both have the same width
+    paddle.bfloat16: np.float16,
+    paddle.float16: np.float16,
+    paddle.int16: np.int16,
+    paddle.uint8: np.uint8,
+    paddle.int8: np.int8,
+    paddle.bool: bool,
+    paddle.float64: np.float64,
+    # XXX: This is ok because both have the same width and byteswap is a no-op anyway
+    paddle.float8_e4m3fn: np.uint8,
+    paddle.float8_e5m2: np.uint8,
+}
+
+
+def _getdtype(dtype_str: str) -> paddle.dtype:
+    return _TYPES[dtype_str]
+
+
+def _view2paddle(safeview, device) -> Dict[str, paddle.Tensor]:
+    result = {}
+    for k, v in safeview:
+        dtype = _getdtype(v["dtype"])
+        if len(v["data"]) == 0:
+            # Workaround because frombuffer doesn't accept zero-size tensors
+            assert any(x == 0 for x in v["shape"])
+            arr = paddle.empty(v["shape"], dtype=dtype)
+        else:
+            arr = paddle.base.core.frombuffer(v["data"], dtype).reshape(v["shape"])
+            if device != "cpu":
+                arr = arr.to(device)
+        if sys.byteorder == "big":
+            arr = paddle.to_tensor(arr.numpy().byteswap(inplace=False), place=device)
+        result[k] = arr
+
+    return result
+
+
+def _tobytes(tensor: paddle.Tensor, name: str) -> bytes:
+    if not tensor.is_contiguous():
+        raise ValueError(
+            f"You are trying to save a non contiguous tensor: `{name}` which is not allowed. It either means you"
+            " are trying to save tensors which are reference of each other in which case it's recommended to save"
+            " only the full tensors, and reslice at load time, or simply call `.contiguous()` on your tensor to"
+            " pack it before saving."
+        )
+    if not tensor.place.is_cpu_place():
+        # Moving tensor to cpu before saving
+        tensor = tensor.cpu()
+
+    import ctypes
+
+    import numpy as np
+
+    # When shape is empty (scalar), np.prod returns a float
+    # we need a int for the following calculations
+    length = int(np.prod(tensor.shape).item())
+    bytes_per_item = _SIZE[tensor.dtype]
+
+    total_bytes = length * bytes_per_item
+
+    ptr = tensor.data_ptr()
+    if ptr == 0:
+        return b""
+    newptr = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_ubyte))
+    data = np.ctypeslib.as_array(newptr, (total_bytes,))  # no internal copy
+    if sys.byteorder == "big":
+        npdtype = NPDTYPES[tensor.dtype]
+        # Not in place as that would potentially modify a live running model
+        data = data.view(npdtype).byteswap(inplace=False)
+    return data.tobytes()
+
+
+def _flatten(tensors: Dict[str, paddle.Tensor]) -> Dict[str, Dict[str, Any]]:
+    if not isinstance(tensors, dict):
+        raise ValueError(
+            f"Expected a dict of [str, paddle.Tensor] but received {type(tensors)}"
+        )
+
+    for k, v in tensors.items():
+        if not isinstance(v, paddle.Tensor):
+            raise ValueError(
+                f"Key `{k}` is invalid, expected paddle.Tensor but received {type(v)}"
+            )
+
+    return {
+        k: {
+            "dtype": str(v.dtype).split(".")[-1],
+            "shape": v.shape,
+            "data": _tobytes(v, k),
+        }
+        for k, v in tensors.items()
+    }

safetensors/torch.py

@@ -221,68 +221,23 @@ def load_model(
     to_removes = _remove_duplicate_names(
         model_state_dict, preferred_names=state_dict.keys()
     )
-
-    reverse_to_remove = {}
-    for key, to_remove_group in to_removes.items():
-        for to_remove in to_remove_group:
-            reverse_to_remove[to_remove] = key
-
-    # We iterate on the model, so we'll add keys we find missing
-    # here
-    missing = set()
-    # We start with all keys on disk declared as unexpected, we'll
-    # slowly remove them when we find them
-    unexpected = set(state_dict.keys())
-    # Some keys can be invalid too.
-    invalid = set()
-
-    for k, mv in model_state_dict.items():
-        actual_k = reverse_to_remove.get(k, None)
-        if actual_k is not None:
-            look_k = actual_k
-        else:
-            look_k = k
-        v = state_dict.get(look_k, None)
-        if v is None:
-            missing.add(k)
-        else:
-            # We can actually check for the shapes while we're at it.
-            # For the device, it's trickier given torch's internals
-            # There might be some Meta device for faster initiation
-            if v.dtype != mv.dtype or v.shape != mv.shape:
-                invalid.add(k)
-            if actual_k is None:
-                unexpected.remove(k)
-
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
     missing = set(missing)
-    unexpected = set(unexpected)
-    if strict and (missing or unexpected or invalid):
+    for to_remove_group in to_removes.values():
+        for to_remove in to_remove_group:
+            if to_remove not in missing:
+                unexpected.append(to_remove)
+            else:
+                missing.remove(to_remove)
+    if strict and (missing or unexpected):
         missing_keys = ", ".join([f'"{k}"' for k in sorted(missing)])
         unexpected_keys = ", ".join([f'"{k}"' for k in sorted(unexpected)])
-        invalid_keys = ", ".join([f'"{k}"' for k in sorted(invalid)])
         error = f"Error(s) in loading state_dict for {model.__class__.__name__}:"
         if missing:
             error += f"\n    Missing key(s) in state_dict: {missing_keys}"
         if unexpected:
             error += f"\n    Unexpected key(s) in state_dict: {unexpected_keys}"
-        if invalid:
-            error += f"\n    Invalid key(s) in state_dict: {invalid_keys}, mismatched dtypes or shape."
-        del state_dict
         raise RuntimeError(error)
-
-    torch_missing, torch_unexpected = model.load_state_dict(state_dict, strict=False)
-    # Sanity check that the work we've done matches
-    # Pytorch internal loading.
-    torch_missing = set(torch_missing)
-    torch_unexpected = set(torch_unexpected)
-    for to_remove_group in to_removes.values():
-        for to_remove in to_remove_group:
-            if to_remove not in torch_missing:
-                torch_unexpected.add(to_remove)
-            else:
-                torch_missing.remove(to_remove)
-    assert torch_missing == missing, f"{torch_missing} != {missing}"
-    assert torch_unexpected == unexpected, f"{torch_unexpected} != {unexpected}"
     return missing, unexpected
 
 
@@ -428,6 +383,7 @@ _SIZE = {
     torch.int8: 1,
     torch.bool: 1,
     torch.float64: 8,
+    torch.complex64: 8,
     _float8_e4m3fn: 1,
     _float8_e5m2: 1,
     _float8_e8m0: 1,
@@ -455,6 +411,7 @@ _TYPES = {
     "BOOL": torch.bool,
     "F8_E4M3": _float8_e4m3fn,
     "F8_E5M2": _float8_e5m2,
+    "C64": torch.complex64,
 }
 if Version(torch.__version__) >= Version("2.3.0"):
     _TYPES.update(
@@ -538,6 +495,7 @@ def _tobytes(tensor: torch.Tensor, name: str) -> bytes:
             # XXX: This is ok because both have the same width and byteswap is a no-op anyway
             _float8_e4m3fn: np.uint8,
             _float8_e5m2: np.uint8,
+            torch.complex64: np.complex64,
         }
         npdtype = NPDTYPES[tensor.dtype]
         # Not in place as that would potentially modify a live running model

{safetensors-0.6.2.dist-info → safetensors-0.7.0rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: safetensors
-Version: 0.6.2
+Version: 0.7.0rc0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Education
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Typing :: Typed
 Requires-Dist: numpy>=1.21.6 ; extra == 'numpy'
+Requires-Dist: packaging ; extra == 'torch'
 Requires-Dist: safetensors[numpy] ; extra == 'torch'
 Requires-Dist: torch>=1.10 ; extra == 'torch'
 Requires-Dist: safetensors[numpy] ; extra == 'tensorflow'

safetensors-0.7.0rc0.dist-info/RECORD

@@ -0,0 +1,14 @@
+safetensors-0.7.0rc0.dist-info/METADATA,sha256=Q7QxCj4WdwYSXi6cJQ2xyMMj1Mtiz0NBsR4n7As7WWs,4188
+safetensors-0.7.0rc0.dist-info/WHEEL,sha256=CG8OzNtm0LMpJ2zhrjswlO8N-965OeMLklsQAG-nMvQ,94
+safetensors-0.7.0rc0.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+safetensors/__init__.py,sha256=HYY5VVsb3b-cxhZBwhNx53ZKqSIB4M14nIXLTOAM1Wc,204
+safetensors/__init__.pyi,sha256=tnVaPqYbh8ggFbOZdYKUC4ArqitiWDfrIQt1BNJ377k,4183
+safetensors/_safetensors_rust.pyd,sha256=SGfYnZXx8FAAnnTG6tfkE5Icp9Fvzo_nABy-BV55SZ8,736768
+safetensors/flax.py,sha256=SnuiGojmth0eCFIWoKEvAfh95nZP9uCZ9E-S4NndrbU,3991
+safetensors/mlx.py,sha256=KvfTWusLSx1hSPWQgg99iL-z9VoD6zQ8l4-RAsCe7P8,3990
+safetensors/numpy.py,sha256=8ci56gDXetlYHH1-Nru83auiUVi-Q1P9bKvfsdkLKPw,5215
+safetensors/paddle.py,sha256=EhXpflqrhKr_NFh4jxV9SUnW0B1vcX_KdPdTqcytrDs,9011
+safetensors/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+safetensors/tensorflow.py,sha256=DajI3qkz00Zy2h7jublSAvTaD51QOPdaIgKQIeSiCRs,4042
+safetensors/torch.py,sha256=CLVWgWQdLm_tVzhRPaeihBHt-4iGAtUW5fY2ys3TyMc,19160
+safetensors-0.7.0rc0.dist-info/RECORD,,

{safetensors-0.6.2.dist-info → safetensors-0.7.0rc0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: maturin (1.9.3)
+Generator: maturin (1.9.6)
 Root-Is-Purelib: false
 Tag: cp38-abi3-win_amd64

safetensors-0.6.2.dist-info/RECORD

@@ -1,14 +0,0 @@
-safetensors-0.6.2.dist-info/METADATA,sha256=kkR-LvpyTCvbdGfZ-pH_IiPF4JNP81N9GPrYjY41BDo,4141
-safetensors-0.6.2.dist-info/WHEEL,sha256=lvaVdaNOIbpDjZxhxQcXMmDSpIrmQUI6MiaH-nloUu8,94
-safetensors-0.6.2.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-safetensors/__init__.py,sha256=HYY5VVsb3b-cxhZBwhNx53ZKqSIB4M14nIXLTOAM1Wc,204
-safetensors/__init__.pyi,sha256=tnVaPqYbh8ggFbOZdYKUC4ArqitiWDfrIQt1BNJ377k,4183
-safetensors/_safetensors_rust.pyd,sha256=05luNkOSehiTNRXXrg4CjBd644FvderG0Xw6TZJPrMc,704000
-safetensors/flax.py,sha256=SnuiGojmth0eCFIWoKEvAfh95nZP9uCZ9E-S4NndrbU,3991
-safetensors/mlx.py,sha256=KvfTWusLSx1hSPWQgg99iL-z9VoD6zQ8l4-RAsCe7P8,3990
-safetensors/numpy.py,sha256=MaUhU4V3J4nDjLmoy0OdeTc6JC8Dq2PBHmnxjQU2bfQ,5189
-safetensors/paddle.py,sha256=B8TLF5MFeqeipUxynSOB_NqPNypL_dE1C3vihtEWj0A,4337
-safetensors/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-safetensors/tensorflow.py,sha256=DajI3qkz00Zy2h7jublSAvTaD51QOPdaIgKQIeSiCRs,4042
-safetensors/torch.py,sha256=wxR0dOwmZBHxv5R5WTmSH4NLOXmvH9JGeVD_uImDlew,20851
-safetensors-0.6.2.dist-info/RECORD,,