safetensors 0.6.1rc0__cp38-abi3-musllinux_1_2_aarch64.whl → 0.7.0rc0__cp38-abi3-musllinux_1_2_aarch64.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,12 +383,13 @@ _SIZE = {
     torch.int8: 1,
     torch.bool: 1,
     torch.float64: 8,
+    torch.complex64: 8,
     _float8_e4m3fn: 1,
     _float8_e5m2: 1,
     _float8_e8m0: 1,
     _float4_e2m1_x2: 1,
 }
-if Version(torch.__version__) > Version("2.0.0"):
+if Version(torch.__version__) >= Version("2.3.0"):
     _SIZE.update(
         {
             torch.uint64: 8,
@@ -455,8 +411,9 @@ _TYPES = {
     "BOOL": torch.bool,
     "F8_E4M3": _float8_e4m3fn,
     "F8_E5M2": _float8_e5m2,
+    "C64": torch.complex64,
 }
-if Version(torch.__version__) > Version("2.0.0"):
+if Version(torch.__version__) >= Version("2.3.0"):
     _TYPES.update(
         {
             "U64": torch.uint64,
@@ -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.1rc0.dist-info → safetensors-0.7.0rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: safetensors
-Version: 0.6.1rc0
+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,15 @@
+safetensors-0.7.0rc0.dist-info/METADATA,sha256=8rEIRH-TlnXY3JsPkaAXv15-RRsnDq7e4cjpAOA7-bQ,4128
+safetensors-0.7.0rc0.dist-info/WHEEL,sha256=bITaocPl86jMm-S2cO-6dnFV3lXGZ1OL6ueZKGDGDMY,106
+safetensors-0.7.0rc0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+safetensors.libs/libgcc_s-39080030.so.1,sha256=fIO6GHOh8Ft9CR0Geu7wSUb9Xnl122iTtrxQQ9TAkTQ,789673
+safetensors/__init__.py,sha256=wbzKZlAVgnAsjHmqryuSJCiADvpDZxNGCfj8VzY0At0,194
+safetensors/__init__.pyi,sha256=IpwsrzRWJA2yR8TxMEC3RHgM_5TiDgSFqyvrAxAa15U,4019
+safetensors/_safetensors_rust.abi3.so,sha256=yYX5CfWL6d_SOU6qHeQ85ZOriLwzovLh9vp2NTakyBU,1379545
+safetensors/flax.py,sha256=T59elUqzVDyGYGdR78QzNEuwyAc8KrEO0EuLBSKOnUs,3853
+safetensors/mlx.py,sha256=IR51jRpcJq6epb0Agj8VsxI9xqBS6NjeAJnr-Ny0jJU,3850
+safetensors/numpy.py,sha256=rit_12-IfZtRgip_VLd8nPAcCXyeM2fPrCDZ7OiyxSY,5028
+safetensors/paddle.py,sha256=LrDwqQbwFnQXiY3M601IU7G6FBctX6tyHHK3_UH6lxE,8721
+safetensors/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+safetensors/tensorflow.py,sha256=AZ-O7-gM-JqTfjczZyCUAHm3Er-GSQnQWaFyY7mAIQc,3903
+safetensors/torch.py,sha256=U0acZVahLsxvqPa1GitMRiaByu6XVbmiBITGtuznBEY,18610
+safetensors-0.7.0rc0.dist-info/RECORD,,

{safetensors-0.6.1rc0.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-musllinux_1_2_aarch64

safetensors-0.6.1rc0.dist-info/RECORD

@@ -1,15 +0,0 @@
-safetensors-0.6.1rc0.dist-info/METADATA,sha256=fifonzsrPqPXgo8sKApMicCCayOiB7N9UU4NFL0t1go,4084
-safetensors-0.6.1rc0.dist-info/WHEEL,sha256=-hAdUMRock4ztuxgCieAciZSiO-8awQp5vVJ2dmPK0M,106
-safetensors-0.6.1rc0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-safetensors.libs/libgcc_s-e52197c3.so.1,sha256=vkPW1Auw6CH9Bjk7frmX3hry_1H9c0tRI0Ncyg71WUI,724137
-safetensors/__init__.py,sha256=wbzKZlAVgnAsjHmqryuSJCiADvpDZxNGCfj8VzY0At0,194
-safetensors/__init__.pyi,sha256=IpwsrzRWJA2yR8TxMEC3RHgM_5TiDgSFqyvrAxAa15U,4019
-safetensors/_safetensors_rust.abi3.so,sha256=ftL5n2W5_mjpt3jUAFlZrl2IO-nFnxFNFxUzV2gHY7k,1248449
-safetensors/flax.py,sha256=T59elUqzVDyGYGdR78QzNEuwyAc8KrEO0EuLBSKOnUs,3853
-safetensors/mlx.py,sha256=IR51jRpcJq6epb0Agj8VsxI9xqBS6NjeAJnr-Ny0jJU,3850
-safetensors/numpy.py,sha256=pIVTnGxAYUSzktkr6xop8B1JK2t3zBH29Ws-oDzo2hA,5003
-safetensors/paddle.py,sha256=0IZqQi_VGKE5Vx179aBOfeD0UBsEGcrndPWvqWSSEG4,4193
-safetensors/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-safetensors/tensorflow.py,sha256=AZ-O7-gM-JqTfjczZyCUAHm3Er-GSQnQWaFyY7mAIQc,3903
-safetensors/torch.py,sha256=3KpStg0fBHGZ4Ii75SxClDIg9Dzu3210dwpOvngUht4,20257
-safetensors-0.6.1rc0.dist-info/RECORD,,