safetensors 0.4.6.dev0__cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl

safetensors/__init__.py

@@ -0,0 +1,9 @@
+# Re-export this
+from ._safetensors_rust import (  # noqa: F401
+    SafetensorError,
+    __version__,
+    deserialize,
+    safe_open,
+    serialize,
+    serialize_file,
+)

safetensors/__init__.pyi

@@ -0,0 +1,73 @@
+# Generated content DO NOT EDIT
+@staticmethod
+def deserialize(bytes):
+    """
+    Opens a safetensors lazily and returns tensors as asked
+
+    Args:
+        data (:obj:`bytes`):
+            The byte content of a file
+
+    Returns:
+        (:obj:`List[str, Dict[str, Dict[str, any]]]`):
+            The deserialized content is like:
+                [("tensor_name", {"shape": [2, 3], "dtype": "F32", "data": b"\0\0.." }), (...)]
+    """
+    pass
+
+@staticmethod
+def serialize(tensor_dict, metadata=None):
+    """
+    Serializes raw data.
+
+    Args:
+        tensor_dict (:obj:`Dict[str, Dict[Any]]`):
+            The tensor dict is like:
+                {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}}
+        metadata (:obj:`Dict[str, str]`, *optional*):
+            The optional purely text annotations
+
+    Returns:
+        (:obj:`bytes`):
+            The serialized content.
+    """
+    pass
+
+@staticmethod
+def serialize_file(tensor_dict, filename, metadata=None):
+    """
+    Serializes raw data.
+
+    Args:
+        tensor_dict (:obj:`Dict[str, Dict[Any]]`):
+            The tensor dict is like:
+                {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}}
+        filename (:obj:`str`):
+            The name of the file to write into.
+        metadata (:obj:`Dict[str, str]`, *optional*):
+            The optional purely text annotations
+
+    Returns:
+        (:obj:`bytes`):
+            The serialized content.
+    """
+    pass
+
+class safe_open:
+    """
+    Opens a safetensors lazily and returns tensors as asked
+
+    Args:
+        filename (:obj:`str`):
+            The filename to open
+
+        framework (:obj:`str`):
+            The framework you want your tensors in. Supported values:
+            `pt`, `tf`, `flax`, `numpy`.
+
+        device (:obj:`str`, defaults to :obj:`"cpu"`):
+            The device on which you want the tensors.
+    """
+
+    def __init__(self, filename, framework, device="cpu"):
+        pass

safetensors/flax.py

@@ -0,0 +1,138 @@
+import os
+from typing import Dict, Optional, Union
+
+import numpy as np
+
+import jax.numpy as jnp
+from jax import Array
+from safetensors import numpy, safe_open
+
+
+def save(tensors: Dict[str, Array], metadata: Optional[Dict[str, str]] = None) -> bytes:
+    """
+    Saves a dictionary of tensors into raw bytes in safetensors format.
+
+    Args:
+        tensors (`Dict[str, Array]`):
+            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.flax import save
+    from jax import numpy as jnp
+
+    tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))}
+    byte_data = save(tensors)
+    ```
+    """
+    np_tensors = _jnp2np(tensors)
+    return numpy.save(np_tensors, metadata=metadata)
+
+
+def save_file(
+    tensors: Dict[str, Array],
+    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, Array]`):
+            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.flax import save_file
+    from jax import numpy as jnp
+
+    tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))}
+    save_file(tensors, "model.safetensors")
+    ```
+    """
+    np_tensors = _jnp2np(tensors)
+    return numpy.save_file(np_tensors, filename, metadata=metadata)
+
+
+def load(data: bytes) -> Dict[str, Array]:
+    """
+    Loads a safetensors file into flax format from pure bytes.
+
+    Args:
+        data (`bytes`):
+            The content of a safetensors file
+
+    Returns:
+        `Dict[str, Array]`: dictionary that contains name as key, value as `Array` on cpu
+
+    Example:
+
+    ```python
+    from safetensors.flax import load
+
+    file_path = "./my_folder/bert.safetensors"
+    with open(file_path, "rb") as f:
+        data = f.read()
+
+    loaded = load(data)
+    ```
+    """
+    flat = numpy.load(data)
+    return _np2jnp(flat)
+
+
+def load_file(filename: Union[str, os.PathLike]) -> Dict[str, Array]:
+    """
+    Loads a safetensors file into flax format.
+
+    Args:
+        filename (`str`, or `os.PathLike`)):
+            The name of the file which contains the tensors
+
+    Returns:
+        `Dict[str, Array]`: dictionary that contains name as key, value as `Array`
+
+    Example:
+
+    ```python
+    from safetensors.flax import load_file
+
+    file_path = "./my_folder/bert.safetensors"
+    loaded = load_file(file_path)
+    ```
+    """
+    result = {}
+    with safe_open(filename, framework="flax") as f:
+        for k in f.keys():
+            result[k] = f.get_tensor(k)
+    return result
+
+
+def _np2jnp(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, Array]:
+    for k, v in numpy_dict.items():
+        numpy_dict[k] = jnp.array(v)
+    return numpy_dict
+
+
+def _jnp2np(jnp_dict: Dict[str, Array]) -> Dict[str, np.array]:
+    for k, v in jnp_dict.items():
+        jnp_dict[k] = np.asarray(v)
+    return jnp_dict

safetensors/mlx.py

@@ -0,0 +1,138 @@
+import os
+from typing import Dict, Optional, Union
+
+import numpy as np
+
+import mlx.core as mx
+from safetensors import numpy, safe_open
+
+
+def save(tensors: Dict[str, mx.array], metadata: Optional[Dict[str, str]] = None) -> bytes:
+    """
+    Saves a dictionary of tensors into raw bytes in safetensors format.
+
+    Args:
+        tensors (`Dict[str, mx.array]`):
+            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.mlx import save
+    import mlx.core as mx
+
+    tensors = {"embedding": mx.zeros((512, 1024)), "attention": mx.zeros((256, 256))}
+    byte_data = save(tensors)
+    ```
+    """
+    np_tensors = _mx2np(tensors)
+    return numpy.save(np_tensors, metadata=metadata)
+
+
+def save_file(
+    tensors: Dict[str, mx.array],
+    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, mx.array]`):
+            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.mlx import save_file
+    import mlx.core as mx
+
+    tensors = {"embedding": mx.zeros((512, 1024)), "attention": mx.zeros((256, 256))}
+    save_file(tensors, "model.safetensors")
+    ```
+    """
+    np_tensors = _mx2np(tensors)
+    return numpy.save_file(np_tensors, filename, metadata=metadata)
+
+
+def load(data: bytes) -> Dict[str, mx.array]:
+    """
+    Loads a safetensors file into MLX format from pure bytes.
+
+    Args:
+        data (`bytes`):
+            The content of a safetensors file
+
+    Returns:
+        `Dict[str, mx.array]`: dictionary that contains name as key, value as `mx.array`
+
+    Example:
+
+    ```python
+    from safetensors.mlx import load
+
+    file_path = "./my_folder/bert.safetensors"
+    with open(file_path, "rb") as f:
+        data = f.read()
+
+    loaded = load(data)
+    ```
+    """
+    flat = numpy.load(data)
+    return _np2mx(flat)
+
+
+def load_file(filename: Union[str, os.PathLike]) -> Dict[str, mx.array]:
+    """
+    Loads a safetensors file into MLX format.
+
+    Args:
+        filename (`str`, or `os.PathLike`)):
+            The name of the file which contains the tensors
+
+    Returns:
+        `Dict[str, mx.array]`: dictionary that contains name as key, value as `mx.array`
+
+    Example:
+
+    ```python
+    from safetensors.flax import load_file
+
+    file_path = "./my_folder/bert.safetensors"
+    loaded = load_file(file_path)
+    ```
+    """
+    result = {}
+    with safe_open(filename, framework="mlx") as f:
+        for k in f.keys():
+            result[k] = f.get_tensor(k)
+    return result
+
+
+def _np2mx(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, mx.array]:
+    for k, v in numpy_dict.items():
+        numpy_dict[k] = mx.array(v)
+    return numpy_dict
+
+
+def _mx2np(mx_dict: Dict[str, mx.array]) -> Dict[str, np.array]:
+    new_dict = {}
+    for k, v in mx_dict.items():
+        new_dict[k] = np.asarray(v)
+    return new_dict

safetensors/numpy.py

@@ -0,0 +1,176 @@
+import os
+import sys
+from typing import Dict, Optional, Union
+
+import numpy as np
+
+from safetensors import deserialize, safe_open, serialize, serialize_file
+
+
+def _tobytes(tensor: np.ndarray) -> bytes:
+    if not _is_little_endian(tensor):
+        tensor = tensor.byteswap(inplace=False)
+    return tensor.tobytes()
+
+
+def save(tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]] = None) -> bytes:
+    """
+    Saves a dictionary of tensors into raw bytes in safetensors format.
+
+    Args:
+        tensor_dict (`Dict[str, np.ndarray]`):
+            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.numpy import save
+    import numpy as np
+
+    tensors = {"embedding": np.zeros((512, 1024)), "attention": np.zeros((256, 256))}
+    byte_data = save(tensors)
+    ```
+    """
+    flattened = {k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)} for k, v in tensor_dict.items()}
+    serialized = serialize(flattened, metadata=metadata)
+    result = bytes(serialized)
+    return result
+
+
+def save_file(
+    tensor_dict: Dict[str, np.ndarray], filename: Union[str, os.PathLike], metadata: Optional[Dict[str, str]] = None
+) -> None:
+    """
+    Saves a dictionary of tensors into raw bytes in safetensors format.
+
+    Args:
+        tensor_dict (`Dict[str, np.ndarray]`):
+            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.numpy import save_file
+    import numpy as np
+
+    tensors = {"embedding": np.zeros((512, 1024)), "attention": np.zeros((256, 256))}
+    save_file(tensors, "model.safetensors")
+    ```
+    """
+    flattened = {k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)} for k, v in tensor_dict.items()}
+    serialize_file(flattened, filename, metadata=metadata)
+
+
+def load(data: bytes) -> Dict[str, np.ndarray]:
+    """
+    Loads a safetensors file into numpy format from pure bytes.
+
+    Args:
+        data (`bytes`):
+            The content of a safetensors file
+
+    Returns:
+        `Dict[str, np.ndarray]`: dictionary that contains name as key, value as `np.ndarray` on cpu
+
+    Example:
+
+    ```python
+    from safetensors.numpy import load
+
+    file_path = "./my_folder/bert.safetensors"
+    with open(file_path, "rb") as f:
+        data = f.read()
+
+    loaded = load(data)
+    ```
+    """
+    flat = deserialize(data)
+    return _view2np(flat)
+
+
+def load_file(filename: Union[str, os.PathLike]) -> Dict[str, np.ndarray]:
+    """
+    Loads a safetensors file into numpy format.
+
+    Args:
+        filename (`str`, or `os.PathLike`)):
+            The name of the file which contains the tensors
+
+    Returns:
+        `Dict[str, np.ndarray]`: dictionary that contains name as key, value as `np.ndarray`
+
+    Example:
+
+    ```python
+    from safetensors.numpy import load_file
+
+    file_path = "./my_folder/bert.safetensors"
+    loaded = load_file(file_path)
+    ```
+    """
+    result = {}
+    with safe_open(filename, framework="np") as f:
+        for k in f.keys():
+            result[k] = f.get_tensor(k)
+    return result
+
+
+_TYPES = {
+    "F64": np.float64,
+    "F32": np.float32,
+    "F16": np.float16,
+    "I64": np.int64,
+    "U64": np.uint64,
+    "I32": np.int32,
+    "U32": np.uint32,
+    "I16": np.int16,
+    "U16": np.uint16,
+    "I8": np.int8,
+    "U8": np.uint8,
+    "BOOL": bool,
+}
+
+
+def _getdtype(dtype_str: str) -> np.dtype:
+    return _TYPES[dtype_str]
+
+
+def _view2np(safeview) -> Dict[str, np.ndarray]:
+    result = {}
+    for k, v in safeview:
+        dtype = _getdtype(v["dtype"])
+        arr = np.frombuffer(v["data"], dtype=dtype).reshape(v["shape"])
+        result[k] = arr
+    return result
+
+
+def _is_little_endian(tensor: np.ndarray) -> bool:
+    byteorder = tensor.dtype.byteorder
+    if byteorder == "=":
+        if sys.byteorder == "little":
+            return True
+        else:
+            return False
+    elif byteorder == "|":
+        return True
+    elif byteorder == "<":
+        return True
+    elif byteorder == ">":
+        return False
+    raise ValueError(f"Unexpected byte order {byteorder}")

safetensors/paddle.py

@@ -0,0 +1,138 @@
+import os
+from typing import Dict, Optional, Union
+
+import numpy as np
+
+import paddle
+from safetensors import numpy
+
+
+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)
+    ```
+    """
+    np_tensors = _paddle2np(tensors)
+    return numpy.save(np_tensors, metadata=metadata)
+
+
+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")
+    ```
+    """
+    np_tensors = _paddle2np(tensors)
+    return numpy.save_file(np_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)
+    ```
+    """
+    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)
+    ```
+    """
+    flat = numpy.load_file(filename)
+    output = _np2paddle(flat, device)
+    return output
+
+
+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