safetensors 0.6.2__tar.gz → 0.7.0rc0__tar.gz

{safetensors-0.6.2 → safetensors-0.7.0rc0}/PKG-INFO

@@ -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.6.2 → safetensors-0.7.0rc0}/bindings/python/Cargo.lock

@@ -2,12 +2,42 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "allocator-api2"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
+
 [[package]]
 name = "autocfg"
 version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
 
+[[package]]
+name = "equivalent"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
+
+[[package]]
+name = "foldhash"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
+
+[[package]]
+name = "hashbrown"
+version = "0.16.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d"
+dependencies = [
+ "allocator-api2",
+ "equivalent",
+ "foldhash",
+ "serde",
+]
+
 [[package]]
 name = "heck"
 version = "0.5.0"
@@ -16,9 +46,12 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
 
 [[package]]
 name = "indoc"
-version = "2.0.6"
+version = "2.0.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd"
+checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
+dependencies = [
+ "rustversion",
+]
 
 [[package]]
 name = "itoa"
@@ -28,21 +61,21 @@ checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
 
 [[package]]
 name = "libc"
-version = "0.2.174"
+version = "0.2.177"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776"
+checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976"
 
 [[package]]
 name = "memchr"
-version = "2.7.5"
+version = "2.7.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0"
+checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273"
 
 [[package]]
 name = "memmap2"
-version = "0.9.5"
+version = "0.9.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f"
+checksum = "744133e4a0e0a658e1374cf3bf8e415c4052a15a111acd372764c55b4177d490"
 dependencies = [
  "libc",
 ]
@@ -70,9 +103,9 @@ checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483"
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.95"
+version = "1.0.103"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "02b3e5e68a3a1a02aad3ec490a98007cbc13c37cbe84a3cd7b8e406d76e7f778"
+checksum = "5ee95bc4ef87b8d5ba32e8b7714ccc834865276eab0aed5c9958d00ec45f49e8"
 dependencies = [
  "unicode-ident",
 ]
@@ -141,13 +174,19 @@ dependencies = [
 
 [[package]]
 name = "quote"
-version = "1.0.40"
+version = "1.0.41"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d"
+checksum = "ce25767e7b499d1b604768e7cde645d14cc8584231ea6b295e9c9eb22c02e1d1"
 dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "rustversion"
+version = "1.0.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
+
 [[package]]
 name = "ryu"
 version = "1.0.20"
@@ -156,15 +195,16 @@ checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
 
 [[package]]
 name = "safetensors"
-version = "0.6.2"
+version = "0.7.0-rc.0"
 dependencies = [
+ "hashbrown",
  "serde",
  "serde_json",
 ]
 
 [[package]]
 name = "safetensors-python"
-version = "0.6.2"
+version = "0.7.0-rc.0"
 dependencies = [
  "memmap2",
  "pyo3",
@@ -174,18 +214,28 @@ dependencies = [
 
 [[package]]
 name = "serde"
-version = "1.0.219"
+version = "1.0.228"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6"
+checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
+dependencies = [
+ "serde_core",
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_core"
+version = "1.0.228"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
 dependencies = [
  "serde_derive",
 ]
 
 [[package]]
 name = "serde_derive"
-version = "1.0.219"
+version = "1.0.228"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00"
+checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -194,21 +244,22 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.140"
+version = "1.0.145"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373"
+checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c"
 dependencies = [
  "itoa",
  "memchr",
  "ryu",
  "serde",
+ "serde_core",
 ]
 
 [[package]]
 name = "syn"
-version = "2.0.104"
+version = "2.0.108"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "17b6f705963418cdb9927482fa304bc562ece2fdd4f616084c50b7023b435a40"
+checksum = "da58917d35242480a05c2897064da0a80589a2a0476c9a3f2fdc83b53502e917"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -217,15 +268,15 @@ dependencies = [
 
 [[package]]
 name = "target-lexicon"
-version = "0.13.2"
+version = "0.13.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a"
+checksum = "df7f62577c25e07834649fc3b39fafdc597c0a3527dc1c60129201ccfcbaa50c"
 
 [[package]]
 name = "unicode-ident"
-version = "1.0.18"
+version = "1.0.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
+checksum = "462eeb75aeb73aea900253ce739c8e18a67423fadf006037cd3ff27e82748a06"
 
 [[package]]
 name = "unindent"

{safetensors-0.6.2 → safetensors-0.7.0rc0}/bindings/python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "safetensors-python"
-version = "0.6.2"
+version = "0.7.0-rc.0"
 edition = "2021"
 rust-version = "1.74"
 readme = "README.md"

{safetensors-0.6.2 → safetensors-0.7.0rc0}/bindings/python/py_src/safetensors/numpy.py

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

safetensors-0.7.0rc0/bindings/python/py_src/safetensors/paddle.py

@@ -0,0 +1,290 @@
+import os
+import sys
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import paddle
+
+from safetensors import numpy, deserialize, safe_open, serialize, serialize_file
+
+
+def save(
+    tensors: Dict[str, paddle.Tensor], metadata: Optional[Dict[str, str]] = None
+) -> bytes:
+    """
+    Saves a dictionary of tensors into raw bytes in safetensors format.
+
+    Args:
+        tensors (`Dict[str, paddle.Tensor]`):
+            The incoming tensors. Tensors need to be contiguous and dense.
+        metadata (`Dict[str, str]`, *optional*, defaults to `None`):
+            Optional text only metadata you might want to save in your header.
+            For instance it can be useful to specify more about the underlying
+            tensors. This is purely informative and does not affect tensor loading.
+
+    Returns:
+        `bytes`: The raw bytes representing the format
+
+    Example:
+
+    ```python
+    from safetensors.paddle import save
+    import paddle
+
+    tensors = {"embedding": paddle.zeros((512, 1024)), "attention": paddle.zeros((256, 256))}
+    byte_data = save(tensors)
+    ```
+    """
+    serialized = serialize(_flatten(tensors), metadata=metadata)
+    result = bytes(serialized)
+    return result
+
+
+def save_file(
+    tensors: Dict[str, paddle.Tensor],
+    filename: Union[str, os.PathLike],
+    metadata: Optional[Dict[str, str]] = None,
+) -> None:
+    """
+    Saves a dictionary of tensors into raw bytes in safetensors format.
+
+    Args:
+        tensors (`Dict[str, paddle.Tensor]`):
+            The incoming tensors. Tensors need to be contiguous and dense.
+        filename (`str`, or `os.PathLike`)):
+            The filename we're saving into.
+        metadata (`Dict[str, str]`, *optional*, defaults to `None`):
+            Optional text only metadata you might want to save in your header.
+            For instance it can be useful to specify more about the underlying
+            tensors. This is purely informative and does not affect tensor loading.
+
+    Returns:
+        `None`
+
+    Example:
+
+    ```python
+    from safetensors.paddle import save_file
+    import paddle
+
+    tensors = {"embedding": paddle.zeros((512, 1024)), "attention": paddle.zeros((256, 256))}
+    save_file(tensors, "model.safetensors")
+    ```
+    """
+    serialize_file(_flatten(tensors), filename, metadata=metadata)
+
+
+def load(data: bytes, device: str = "cpu") -> Dict[str, paddle.Tensor]:
+    """
+    Loads a safetensors file into paddle format from pure bytes.
+
+    Args:
+        data (`bytes`):
+            The content of a safetensors file
+
+    Returns:
+        `Dict[str, paddle.Tensor]`: dictionary that contains name as key, value as `paddle.Tensor` on cpu
+
+    Example:
+
+    ```python
+    from safetensors.paddle import load
+
+    file_path = "./my_folder/bert.safetensors"
+    with open(file_path, "rb") as f:
+        data = f.read()
+
+    loaded = load(data)
+    ```
+    """
+    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(
+    filename: Union[str, os.PathLike], device="cpu"
+) -> Dict[str, paddle.Tensor]:
+    """
+    Loads a safetensors file into paddle format.
+
+    Args:
+        filename (`str`, or `os.PathLike`)):
+            The name of the file which contains the tensors
+        device (`Union[Dict[str, any], str]`, *optional*, defaults to `cpu`):
+            The device where the tensors need to be located after load.
+            available options are all regular paddle device locations
+
+    Returns:
+        `Dict[str, paddle.Tensor]`: dictionary that contains name as key, value as `paddle.Tensor`
+
+    Example:
+
+    ```python
+    from safetensors.paddle import load_file
+
+    file_path = "./my_folder/bert.safetensors"
+    loaded = load_file(file_path)
+    ```
+    """
+    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(
+    numpy_dict: Dict[str, np.ndarray], device: str = "cpu"
+) -> Dict[str, paddle.Tensor]:
+    for k, v in numpy_dict.items():
+        numpy_dict[k] = paddle.to_tensor(v, place=device)
+    return numpy_dict
+
+
+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-0.6.2 → safetensors-0.7.0rc0/bindings/python}/py_src/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