Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,10 @@
|
||||
# Re-export this
|
||||
from ._safetensors_rust import ( # noqa: F401
|
||||
SafetensorError,
|
||||
__version__,
|
||||
deserialize,
|
||||
safe_open,
|
||||
_safe_open_handle,
|
||||
serialize,
|
||||
serialize_file,
|
||||
)
|
||||
@@ -0,0 +1,164 @@
|
||||
# Generated content DO NOT EDIT
|
||||
@staticmethod
|
||||
def deserialize(bytes):
|
||||
"""
|
||||
Opens a safetensors lazily and returns tensors as asked
|
||||
|
||||
Args:
|
||||
data (`bytes`):
|
||||
The byte content of a file
|
||||
|
||||
Returns:
|
||||
(`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 (`Dict[str, Dict[Any]]`):
|
||||
The tensor dict is like:
|
||||
{"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}}
|
||||
metadata (`Dict[str, str]`, *optional*):
|
||||
The optional purely text annotations
|
||||
|
||||
Returns:
|
||||
(`bytes`):
|
||||
The serialized content.
|
||||
"""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def serialize_file(tensor_dict, filename, metadata=None):
|
||||
"""
|
||||
Serializes raw data into file.
|
||||
|
||||
Args:
|
||||
tensor_dict (`Dict[str, Dict[Any]]`):
|
||||
The tensor dict is like:
|
||||
{"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}}
|
||||
filename (`str`, or `os.PathLike`):
|
||||
The name of the file to write into.
|
||||
metadata (`Dict[str, str]`, *optional*):
|
||||
The optional purely text annotations
|
||||
|
||||
Returns:
|
||||
(`NoneType`):
|
||||
On success return None
|
||||
"""
|
||||
pass
|
||||
|
||||
class safe_open:
|
||||
"""
|
||||
Opens a safetensors lazily and returns tensors as asked
|
||||
|
||||
Args:
|
||||
filename (`str`, or `os.PathLike`):
|
||||
The filename to open
|
||||
|
||||
framework (`str`):
|
||||
The framework you want you tensors in. Supported values:
|
||||
`pt`, `tf`, `flax`, `numpy`.
|
||||
|
||||
device (`str`, defaults to `"cpu"`):
|
||||
The device on which you want the tensors.
|
||||
"""
|
||||
def __init__(self, filename, framework, device=...):
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
"""
|
||||
Start the context manager
|
||||
"""
|
||||
pass
|
||||
|
||||
def __exit__(self, _exc_type, _exc_value, _traceback):
|
||||
"""
|
||||
Exits the context manager
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_slice(self, name):
|
||||
"""
|
||||
Returns a full slice view object
|
||||
|
||||
Args:
|
||||
name (`str`):
|
||||
The name of the tensor you want
|
||||
|
||||
Returns:
|
||||
(`PySafeSlice`):
|
||||
A dummy object you can slice into to get a real tensor
|
||||
Example:
|
||||
```python
|
||||
from safetensors import safe_open
|
||||
|
||||
with safe_open("model.safetensors", framework="pt", device=0) as f:
|
||||
tensor_part = f.get_slice("embedding")[:, ::8]
|
||||
|
||||
```
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_tensor(self, name):
|
||||
"""
|
||||
Returns a full tensor
|
||||
|
||||
Args:
|
||||
name (`str`):
|
||||
The name of the tensor you want
|
||||
|
||||
Returns:
|
||||
(`Tensor`):
|
||||
The tensor in the framework you opened the file for.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from safetensors import safe_open
|
||||
|
||||
with safe_open("model.safetensors", framework="pt", device=0) as f:
|
||||
tensor = f.get_tensor("embedding")
|
||||
|
||||
```
|
||||
"""
|
||||
pass
|
||||
|
||||
def keys(self):
|
||||
"""
|
||||
Returns the names of the tensors in the file.
|
||||
|
||||
Returns:
|
||||
(`List[str]`):
|
||||
The name of the tensors contained in that file
|
||||
"""
|
||||
pass
|
||||
|
||||
def metadata(self):
|
||||
"""
|
||||
Return the special non tensor information in the header
|
||||
|
||||
Returns:
|
||||
(`Dict[str, str]`):
|
||||
The freeform metadata.
|
||||
"""
|
||||
pass
|
||||
|
||||
def offset_keys(self):
|
||||
"""
|
||||
Returns the names of the tensors in the file, ordered by offset.
|
||||
|
||||
Returns:
|
||||
(`List[str]`):
|
||||
The name of the tensors contained in that file
|
||||
"""
|
||||
pass
|
||||
|
||||
class SafetensorError(Exception):
|
||||
"""
|
||||
Custom Python Exception for Safetensor errors.
|
||||
"""
|
||||
BIN
Binary file not shown.
@@ -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.offset_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
|
||||
@@ -0,0 +1,140 @@
|
||||
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.offset_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
|
||||
@@ -0,0 +1,187 @@
|
||||
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.offset_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,
|
||||
"C64": np.complex64,
|
||||
}
|
||||
|
||||
|
||||
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}")
|
||||
@@ -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()
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
import os
|
||||
from typing import Dict, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
|
||||
from safetensors import numpy, safe_open
|
||||
|
||||
|
||||
def save(
|
||||
tensors: Dict[str, tf.Tensor], metadata: Optional[Dict[str, str]] = None
|
||||
) -> bytes:
|
||||
"""
|
||||
Saves a dictionary of tensors into raw bytes in safetensors format.
|
||||
|
||||
Args:
|
||||
tensors (`Dict[str, tf.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.tensorflow import save
|
||||
import tensorflow as tf
|
||||
|
||||
tensors = {"embedding": tf.zeros((512, 1024)), "attention": tf.zeros((256, 256))}
|
||||
byte_data = save(tensors)
|
||||
```
|
||||
"""
|
||||
np_tensors = _tf2np(tensors)
|
||||
return numpy.save(np_tensors, metadata=metadata)
|
||||
|
||||
|
||||
def save_file(
|
||||
tensors: Dict[str, tf.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, tf.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.tensorflow import save_file
|
||||
import tensorflow as tf
|
||||
|
||||
tensors = {"embedding": tf.zeros((512, 1024)), "attention": tf.zeros((256, 256))}
|
||||
save_file(tensors, "model.safetensors")
|
||||
```
|
||||
"""
|
||||
np_tensors = _tf2np(tensors)
|
||||
return numpy.save_file(np_tensors, filename, metadata=metadata)
|
||||
|
||||
|
||||
def load(data: bytes) -> Dict[str, tf.Tensor]:
|
||||
"""
|
||||
Loads a safetensors file into tensorflow format from pure bytes.
|
||||
|
||||
Args:
|
||||
data (`bytes`):
|
||||
The content of a safetensors file
|
||||
|
||||
Returns:
|
||||
`Dict[str, tf.Tensor]`: dictionary that contains name as key, value as `tf.Tensor` on cpu
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
from safetensors.tensorflow 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 _np2tf(flat)
|
||||
|
||||
|
||||
def load_file(filename: Union[str, os.PathLike]) -> Dict[str, tf.Tensor]:
|
||||
"""
|
||||
Loads a safetensors file into tensorflow format.
|
||||
|
||||
Args:
|
||||
filename (`str`, or `os.PathLike`)):
|
||||
The name of the file which contains the tensors
|
||||
|
||||
Returns:
|
||||
`Dict[str, tf.Tensor]`: dictionary that contains name as key, value as `tf.Tensor`
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
from safetensors.tensorflow import load_file
|
||||
|
||||
file_path = "./my_folder/bert.safetensors"
|
||||
loaded = load_file(file_path)
|
||||
```
|
||||
"""
|
||||
result = {}
|
||||
with safe_open(filename, framework="tf") as f:
|
||||
for k in f.offset_keys():
|
||||
result[k] = f.get_tensor(k)
|
||||
return result
|
||||
|
||||
|
||||
def _np2tf(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, tf.Tensor]:
|
||||
for k, v in numpy_dict.items():
|
||||
numpy_dict[k] = tf.convert_to_tensor(v)
|
||||
return numpy_dict
|
||||
|
||||
|
||||
def _tf2np(tf_dict: Dict[str, tf.Tensor]) -> Dict[str, np.array]:
|
||||
for k, v in tf_dict.items():
|
||||
tf_dict[k] = v.numpy()
|
||||
return tf_dict
|
||||
@@ -0,0 +1,550 @@
|
||||
import os
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple, Union
|
||||
from packaging.version import Version
|
||||
|
||||
import torch
|
||||
|
||||
from safetensors import deserialize, safe_open, serialize, serialize_file
|
||||
|
||||
|
||||
def storage_ptr(tensor: torch.Tensor) -> int:
|
||||
try:
|
||||
return tensor.untyped_storage().data_ptr()
|
||||
except Exception:
|
||||
# Fallback for torch==1.10
|
||||
try:
|
||||
return tensor.storage().data_ptr()
|
||||
except NotImplementedError:
|
||||
# Fallback for meta storage
|
||||
return 0
|
||||
|
||||
|
||||
def _end_ptr(tensor: torch.Tensor) -> int:
|
||||
if tensor.nelement():
|
||||
stop = tensor.view(-1)[-1].data_ptr() + _SIZE[tensor.dtype]
|
||||
else:
|
||||
stop = tensor.data_ptr()
|
||||
return stop
|
||||
|
||||
|
||||
def storage_size(tensor: torch.Tensor) -> int:
|
||||
try:
|
||||
return tensor.untyped_storage().nbytes()
|
||||
except AttributeError:
|
||||
# Fallback for torch==1.10
|
||||
try:
|
||||
return tensor.storage().size() * _SIZE[tensor.dtype]
|
||||
except NotImplementedError:
|
||||
# Fallback for meta storage
|
||||
# On torch >=2.0 this is the tensor size
|
||||
return tensor.nelement() * _SIZE[tensor.dtype]
|
||||
|
||||
|
||||
def _filter_shared_not_shared(
|
||||
tensors: List[Set[str]], state_dict: Dict[str, torch.Tensor]
|
||||
) -> List[Set[str]]:
|
||||
filtered_tensors = []
|
||||
for shared in tensors:
|
||||
if len(shared) < 2:
|
||||
filtered_tensors.append(shared)
|
||||
continue
|
||||
|
||||
areas = []
|
||||
for name in shared:
|
||||
tensor = state_dict[name]
|
||||
areas.append((tensor.data_ptr(), _end_ptr(tensor), name))
|
||||
areas.sort()
|
||||
|
||||
_, last_stop, last_name = areas[0]
|
||||
filtered_tensors.append({last_name})
|
||||
for start, stop, name in areas[1:]:
|
||||
if start >= last_stop:
|
||||
filtered_tensors.append({name})
|
||||
else:
|
||||
filtered_tensors[-1].add(name)
|
||||
last_stop = stop
|
||||
|
||||
return filtered_tensors
|
||||
|
||||
|
||||
def _find_shared_tensors(state_dict: Dict[str, torch.Tensor]) -> List[Set[str]]:
|
||||
tensors = defaultdict(set)
|
||||
for k, v in state_dict.items():
|
||||
if (
|
||||
v.device != torch.device("meta")
|
||||
and storage_ptr(v) != 0
|
||||
and storage_size(v) != 0
|
||||
):
|
||||
# Need to add device as key because of multiple GPU.
|
||||
tensors[(v.device, storage_ptr(v), storage_size(v))].add(k)
|
||||
tensors = list(sorted(tensors.values()))
|
||||
tensors = _filter_shared_not_shared(tensors, state_dict)
|
||||
return tensors
|
||||
|
||||
|
||||
def _is_complete(tensor: torch.Tensor) -> bool:
|
||||
return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _SIZE[
|
||||
tensor.dtype
|
||||
] == storage_size(tensor)
|
||||
|
||||
|
||||
def _remove_duplicate_names(
|
||||
state_dict: Dict[str, torch.Tensor],
|
||||
*,
|
||||
preferred_names: Optional[List[str]] = None,
|
||||
discard_names: Optional[List[str]] = None,
|
||||
) -> Dict[str, List[str]]:
|
||||
if preferred_names is None:
|
||||
preferred_names = []
|
||||
preferred_names = set(preferred_names)
|
||||
if discard_names is None:
|
||||
discard_names = []
|
||||
discard_names = set(discard_names)
|
||||
|
||||
shareds = _find_shared_tensors(state_dict)
|
||||
to_remove = defaultdict(list)
|
||||
for shared in shareds:
|
||||
complete_names = set(
|
||||
[name for name in shared if _is_complete(state_dict[name])]
|
||||
)
|
||||
if not complete_names:
|
||||
raise RuntimeError(
|
||||
"Error while trying to find names to remove to save state dict, but found no suitable name to keep"
|
||||
f" for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model"
|
||||
" since you could be storing much more memory than needed. Please refer to"
|
||||
" https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an"
|
||||
" issue."
|
||||
)
|
||||
|
||||
keep_name = sorted(list(complete_names))[0]
|
||||
|
||||
# Mechanism to preferentially select keys to keep
|
||||
# coming from the on-disk file to allow
|
||||
# loading models saved with a different choice
|
||||
# of keep_name
|
||||
preferred = complete_names.difference(discard_names)
|
||||
if preferred:
|
||||
keep_name = sorted(list(preferred))[0]
|
||||
|
||||
if preferred_names:
|
||||
preferred = preferred_names.intersection(complete_names)
|
||||
if preferred:
|
||||
keep_name = sorted(list(preferred))[0]
|
||||
for name in sorted(shared):
|
||||
if name != keep_name:
|
||||
to_remove[keep_name].append(name)
|
||||
return to_remove
|
||||
|
||||
|
||||
def save_model(
|
||||
model: torch.nn.Module,
|
||||
filename: str,
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
force_contiguous: bool = True,
|
||||
):
|
||||
"""
|
||||
Saves a given torch model to specified filename.
|
||||
This method exists specifically to avoid tensor sharing issues which are
|
||||
not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors)
|
||||
|
||||
Args:
|
||||
model (`torch.nn.Module`):
|
||||
The model to save on disk.
|
||||
filename (`str`):
|
||||
The filename location to save the file
|
||||
metadata (`Dict[str, str]`, *optional*):
|
||||
Extra information to save along with the file.
|
||||
Some metadata will be added for each dropped tensors.
|
||||
This information will not be enough to recover the entire
|
||||
shared structure but might help understanding things
|
||||
force_contiguous (`boolean`, *optional*, defaults to True):
|
||||
Forcing the state_dict to be saved as contiguous tensors.
|
||||
This has no effect on the correctness of the model, but it
|
||||
could potentially change performance if the layout of the tensor
|
||||
was chosen specifically for that reason.
|
||||
"""
|
||||
state_dict = model.state_dict()
|
||||
to_removes = _remove_duplicate_names(state_dict)
|
||||
|
||||
for kept_name, to_remove_group in to_removes.items():
|
||||
for to_remove in to_remove_group:
|
||||
if metadata is None:
|
||||
metadata = {}
|
||||
|
||||
if to_remove not in metadata:
|
||||
# Do not override user data
|
||||
metadata[to_remove] = kept_name
|
||||
del state_dict[to_remove]
|
||||
if force_contiguous:
|
||||
state_dict = {k: v.contiguous() for k, v in state_dict.items()}
|
||||
try:
|
||||
save_file(state_dict, filename, metadata=metadata)
|
||||
except ValueError as e:
|
||||
msg = str(e)
|
||||
msg += " Or use save_model(..., force_contiguous=True), read the docs for potential caveats."
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def load_model(
|
||||
model: torch.nn.Module,
|
||||
filename: Union[str, os.PathLike],
|
||||
strict: bool = True,
|
||||
device: Union[str, int] = "cpu",
|
||||
) -> Tuple[List[str], List[str]]:
|
||||
"""
|
||||
Loads a given filename onto a torch model.
|
||||
This method exists specifically to avoid tensor sharing issues which are
|
||||
not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors)
|
||||
|
||||
Args:
|
||||
model (`torch.nn.Module`):
|
||||
The model to load onto.
|
||||
filename (`str`, or `os.PathLike`):
|
||||
The filename location to load the file from.
|
||||
strict (`bool`, *optional*, defaults to True):
|
||||
Whether to fail if you're missing keys or having unexpected ones.
|
||||
When false, the function simply returns missing and unexpected names.
|
||||
device (`Union[str, int]`, *optional*, defaults to `cpu`):
|
||||
The device where the tensors need to be located after load.
|
||||
available options are all regular torch device locations.
|
||||
|
||||
Returns:
|
||||
`(missing, unexpected): (List[str], List[str])`
|
||||
`missing` are names in the model which were not modified during loading
|
||||
`unexpected` are names that are on the file, but weren't used during
|
||||
the load.
|
||||
"""
|
||||
state_dict = load_file(filename, device=device)
|
||||
model_state_dict = model.state_dict()
|
||||
to_removes = _remove_duplicate_names(
|
||||
model_state_dict, preferred_names=state_dict.keys()
|
||||
)
|
||||
missing, unexpected = model.load_state_dict(state_dict, strict=False)
|
||||
missing = set(missing)
|
||||
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)])
|
||||
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}"
|
||||
raise RuntimeError(error)
|
||||
return missing, unexpected
|
||||
|
||||
|
||||
def save(
|
||||
tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None
|
||||
) -> bytes:
|
||||
"""
|
||||
Saves a dictionary of tensors into raw bytes in safetensors format.
|
||||
|
||||
Args:
|
||||
tensors (`Dict[str, torch.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.torch import save
|
||||
import torch
|
||||
|
||||
tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
|
||||
byte_data = save(tensors)
|
||||
```
|
||||
"""
|
||||
serialized = serialize(_flatten(tensors), metadata=metadata)
|
||||
result = bytes(serialized)
|
||||
return result
|
||||
|
||||
|
||||
def save_file(
|
||||
tensors: Dict[str, torch.Tensor],
|
||||
filename: Union[str, os.PathLike],
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
):
|
||||
"""
|
||||
Saves a dictionary of tensors into raw bytes in safetensors format.
|
||||
|
||||
Args:
|
||||
tensors (`Dict[str, torch.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.torch import save_file
|
||||
import torch
|
||||
|
||||
tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
|
||||
save_file(tensors, "model.safetensors")
|
||||
```
|
||||
"""
|
||||
serialize_file(_flatten(tensors), filename, metadata=metadata)
|
||||
|
||||
|
||||
def load_file(
|
||||
filename: Union[str, os.PathLike], device: Union[str, int] = "cpu"
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
Loads a safetensors file into torch format.
|
||||
|
||||
Args:
|
||||
filename (`str`, or `os.PathLike`):
|
||||
The name of the file which contains the tensors
|
||||
device (`Union[str, int]`, *optional*, defaults to `cpu`):
|
||||
The device where the tensors need to be located after load.
|
||||
available options are all regular torch device locations.
|
||||
|
||||
Returns:
|
||||
`Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor`
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
from safetensors.torch import load_file
|
||||
|
||||
file_path = "./my_folder/bert.safetensors"
|
||||
loaded = load_file(file_path)
|
||||
```
|
||||
"""
|
||||
result = {}
|
||||
with safe_open(filename, framework="pt", device=device) as f:
|
||||
for k in f.offset_keys():
|
||||
result[k] = f.get_tensor(k)
|
||||
return result
|
||||
|
||||
|
||||
def load(data: bytes) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
Loads a safetensors file into torch format from pure bytes.
|
||||
|
||||
Args:
|
||||
data (`bytes`):
|
||||
The content of a safetensors file
|
||||
|
||||
Returns:
|
||||
`Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor` on cpu
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
from safetensors.torch 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 _view2torch(flat)
|
||||
|
||||
|
||||
# torch.float8 formats require 2.1; we do not support these dtypes on earlier versions
|
||||
_float8_e4m3fn = getattr(torch, "float8_e4m3fn", None)
|
||||
_float8_e5m2 = getattr(torch, "float8_e5m2", None)
|
||||
_float8_e8m0 = getattr(torch, "float8_e8m0fnu", None)
|
||||
_float4_e2m1_x2 = getattr(torch, "float4_e2m1fn_x2", None)
|
||||
|
||||
_SIZE = {
|
||||
torch.int64: 8,
|
||||
torch.float32: 4,
|
||||
torch.int32: 4,
|
||||
torch.bfloat16: 2,
|
||||
torch.float16: 2,
|
||||
torch.int16: 2,
|
||||
torch.uint8: 1,
|
||||
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.3.0"):
|
||||
_SIZE.update(
|
||||
{
|
||||
torch.uint64: 8,
|
||||
torch.uint32: 4,
|
||||
torch.uint16: 2,
|
||||
}
|
||||
)
|
||||
|
||||
_TYPES = {
|
||||
"F64": torch.float64,
|
||||
"F32": torch.float32,
|
||||
"F16": torch.float16,
|
||||
"BF16": torch.bfloat16,
|
||||
"I64": torch.int64,
|
||||
"I32": torch.int32,
|
||||
"I16": torch.int16,
|
||||
"I8": torch.int8,
|
||||
"U8": torch.uint8,
|
||||
"BOOL": torch.bool,
|
||||
"F8_E4M3": _float8_e4m3fn,
|
||||
"F8_E5M2": _float8_e5m2,
|
||||
"C64": torch.complex64,
|
||||
}
|
||||
if Version(torch.__version__) >= Version("2.3.0"):
|
||||
_TYPES.update(
|
||||
{
|
||||
"U64": torch.uint64,
|
||||
"U32": torch.uint32,
|
||||
"U16": torch.uint16,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _getdtype(dtype_str: str) -> torch.dtype:
|
||||
return _TYPES[dtype_str]
|
||||
|
||||
|
||||
def _view2torch(safeview) -> Dict[str, torch.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 = torch.empty(v["shape"], dtype=dtype)
|
||||
else:
|
||||
arr = torch.frombuffer(v["data"], dtype=dtype).reshape(v["shape"])
|
||||
if sys.byteorder == "big":
|
||||
arr = torch.from_numpy(arr.numpy().byteswap(inplace=False))
|
||||
result[k] = arr
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _tobytes(tensor: torch.Tensor, name: str) -> bytes:
|
||||
if tensor.layout != torch.strided:
|
||||
raise ValueError(
|
||||
f"You are trying to save a sparse tensor: `{name}` which this library does not support."
|
||||
" You can make it a dense tensor before saving with `.to_dense()` but be aware this might"
|
||||
" make a much larger file than needed."
|
||||
)
|
||||
|
||||
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 tensor.device.type != "cpu":
|
||||
# Moving tensor to cpu before saving
|
||||
tensor = tensor.to("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":
|
||||
NPDTYPES = {
|
||||
torch.int64: np.int64,
|
||||
torch.float32: np.float32,
|
||||
torch.int32: np.int32,
|
||||
# XXX: This is ok because both have the same width
|
||||
torch.bfloat16: np.float16,
|
||||
torch.float16: np.float16,
|
||||
torch.int16: np.int16,
|
||||
torch.uint8: np.uint8,
|
||||
torch.int8: np.int8,
|
||||
torch.bool: bool,
|
||||
torch.float64: np.float64,
|
||||
# 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
|
||||
data = data.view(npdtype).byteswap(inplace=False)
|
||||
return data.tobytes()
|
||||
|
||||
|
||||
def _flatten(tensors: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, Any]]:
|
||||
if not isinstance(tensors, dict):
|
||||
raise ValueError(
|
||||
f"Expected a dict of [str, torch.Tensor] but received {type(tensors)}"
|
||||
)
|
||||
|
||||
invalid_tensors = []
|
||||
for k, v in tensors.items():
|
||||
if not isinstance(v, torch.Tensor):
|
||||
raise ValueError(
|
||||
f"Key `{k}` is invalid, expected torch.Tensor but received {type(v)}"
|
||||
)
|
||||
|
||||
if v.layout != torch.strided:
|
||||
invalid_tensors.append(k)
|
||||
if invalid_tensors:
|
||||
raise ValueError(
|
||||
f"You are trying to save a sparse tensors: `{invalid_tensors}` which this library does not support."
|
||||
" You can make it a dense tensor before saving with `.to_dense()` but be aware this might"
|
||||
" make a much larger file than needed."
|
||||
)
|
||||
|
||||
shared_pointers = _find_shared_tensors(tensors)
|
||||
failing = []
|
||||
for names in shared_pointers:
|
||||
if len(names) > 1:
|
||||
failing.append(names)
|
||||
|
||||
if failing:
|
||||
raise RuntimeError(
|
||||
f"""
|
||||
Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: {failing}.
|
||||
A potential way to correctly save your model is to use `save_model`.
|
||||
More information at https://huggingface.co/docs/safetensors/torch_shared_tensors
|
||||
"""
|
||||
)
|
||||
|
||||
return {
|
||||
k: {
|
||||
"dtype": str(v.dtype).split(".")[-1],
|
||||
"shape": v.shape,
|
||||
"data": _tobytes(v, k),
|
||||
}
|
||||
for k, v in tensors.items()
|
||||
}
|
||||
Reference in New Issue
Block a user