Upload 23 files

.gitattributes

@@ -511,3 +511,4 @@ python_env/lib/site-packages/scipy/stats/tests/__pycache__/test_distributions.cp
 python_env/lib/site-packages/scipy/stats/tests/__pycache__/test_morestats.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
 python_env/lib/site-packages/scipy/stats/tests/__pycache__/test_multivariate.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
 python_env/lib/site-packages/scipy/stats/tests/__pycache__/test_stats.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
+python_env/lib/site-packages/safetensors/_safetensors_rust.pyd filter=lfs diff=lfs merge=lfs -text

python_env/lib/site-packages/safetensors-0.7.0.dist-info/INSTALLER

@@ -0,0 +1 @@
+pip

python_env/lib/site-packages/safetensors-0.7.0.dist-info/METADATA

@@ -0,0 +1,133 @@
+Metadata-Version: 2.4
+Name: safetensors
+Version: 0.7.0
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+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'
+Requires-Dist: tensorflow>=2.11.0 ; extra == 'tensorflow'
+Requires-Dist: safetensors[numpy] ; extra == 'pinned-tf'
+Requires-Dist: tensorflow==2.18.0 ; extra == 'pinned-tf'
+Requires-Dist: safetensors[numpy] ; extra == 'jax'
+Requires-Dist: flax>=0.6.3 ; extra == 'jax'
+Requires-Dist: jax>=0.3.25 ; extra == 'jax'
+Requires-Dist: jaxlib>=0.3.25 ; extra == 'jax'
+Requires-Dist: mlx>=0.0.9 ; extra == 'mlx'
+Requires-Dist: safetensors[numpy] ; extra == 'paddlepaddle'
+Requires-Dist: paddlepaddle>=2.4.1 ; extra == 'paddlepaddle'
+Requires-Dist: ruff ; extra == 'quality'
+Requires-Dist: safetensors[numpy] ; extra == 'testing'
+Requires-Dist: h5py>=3.7.0 ; extra == 'testing'
+Requires-Dist: huggingface-hub>=0.12.1 ; extra == 'testing'
+Requires-Dist: setuptools-rust>=1.5.2 ; extra == 'testing'
+Requires-Dist: pytest>=7.2.0 ; extra == 'testing'
+Requires-Dist: pytest-benchmark>=4.0.0 ; extra == 'testing'
+Requires-Dist: hypothesis>=6.70.2 ; extra == 'testing'
+Requires-Dist: safetensors[numpy] ; extra == 'testingfree'
+Requires-Dist: huggingface-hub>=0.12.1 ; extra == 'testingfree'
+Requires-Dist: setuptools-rust>=1.5.2 ; extra == 'testingfree'
+Requires-Dist: pytest>=7.2.0 ; extra == 'testingfree'
+Requires-Dist: pytest-benchmark>=4.0.0 ; extra == 'testingfree'
+Requires-Dist: hypothesis>=6.70.2 ; extra == 'testingfree'
+Requires-Dist: safetensors[torch] ; extra == 'all'
+Requires-Dist: safetensors[numpy] ; extra == 'all'
+Requires-Dist: safetensors[pinned-tf] ; extra == 'all'
+Requires-Dist: safetensors[jax] ; extra == 'all'
+Requires-Dist: safetensors[paddlepaddle] ; extra == 'all'
+Requires-Dist: safetensors[quality] ; extra == 'all'
+Requires-Dist: safetensors[testing] ; extra == 'all'
+Requires-Dist: safetensors[all] ; extra == 'dev'
+Provides-Extra: numpy
+Provides-Extra: torch
+Provides-Extra: tensorflow
+Provides-Extra: pinned-tf
+Provides-Extra: jax
+Provides-Extra: mlx
+Provides-Extra: paddlepaddle
+Provides-Extra: quality
+Provides-Extra: testing
+Provides-Extra: testingfree
+Provides-Extra: all
+Provides-Extra: dev
+License-File: LICENSE
+Author-email: Nicolas Patry <patry.nicolas@protonmail.com>
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Homepage, https://github.com/huggingface/safetensors
+Project-URL: Source, https://github.com/huggingface/safetensors
+
+## Installation
+
+```
+pip install safetensors
+```
+
+
+## Usage
+
+### Numpy
+
+```python
+from safetensors.numpy import save_file, load_file
+import numpy as np
+
+tensors = {
+   "a": np.zeros((2, 2)),
+   "b": np.zeros((2, 3), dtype=np.uint8)
+}
+
+save_file(tensors, "./model.safetensors")
+
+
+# Now loading
+loaded = load_file("./model.safetensors")
+```
+
+### Torch
+
+```python
+from safetensors.torch import save_file, load_file
+import torch
+
+tensors = {
+   "a": torch.zeros((2, 2)),
+   "b": torch.zeros((2, 3), dtype=torch.uint8)
+}
+
+save_file(tensors, "./model.safetensors")
+
+
+# Now loading
+loaded = load_file("./model.safetensors")
+```
+
+### Developing
+
+```
+# inside ./safetensors/bindings/python
+pip install .[dev]
+```
+Should be enough to install this library locally.
+
+### Testing
+
+```
+# inside ./safetensors/bindings/python
+pip install .[dev]
+pytest -sv tests/
+```
+

python_env/lib/site-packages/safetensors-0.7.0.dist-info/RECORD

@@ -0,0 +1,23 @@
+safetensors-0.7.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+safetensors-0.7.0.dist-info/METADATA,sha256=j4yXpt97GFoHwriefq-dEwV9ZX_7SZKS795bMtYUiBs,4185
+safetensors-0.7.0.dist-info/RECORD,,
+safetensors-0.7.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+safetensors-0.7.0.dist-info/WHEEL,sha256=MW5GXj0PFGHxdts4kUIyLDU55oFjz5S6uZmcnjsRmds,95
+safetensors-0.7.0.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+safetensors/__init__.py,sha256=HYY5VVsb3b-cxhZBwhNx53ZKqSIB4M14nIXLTOAM1Wc,204
+safetensors/__init__.pyi,sha256=tnVaPqYbh8ggFbOZdYKUC4ArqitiWDfrIQt1BNJ377k,4183
+safetensors/__pycache__/__init__.cpython-310.pyc,,
+safetensors/__pycache__/flax.cpython-310.pyc,,
+safetensors/__pycache__/mlx.cpython-310.pyc,,
+safetensors/__pycache__/numpy.cpython-310.pyc,,
+safetensors/__pycache__/paddle.cpython-310.pyc,,
+safetensors/__pycache__/tensorflow.cpython-310.pyc,,
+safetensors/__pycache__/torch.cpython-310.pyc,,
+safetensors/_safetensors_rust.pyd,sha256=8gPLI70OVyFz6ze8U_qlSH1QD5vCVAkB0NpGbFq62vA,738304
+safetensors/flax.py,sha256=SnuiGojmth0eCFIWoKEvAfh95nZP9uCZ9E-S4NndrbU,3991
+safetensors/mlx.py,sha256=KvfTWusLSx1hSPWQgg99iL-z9VoD6zQ8l4-RAsCe7P8,3990
+safetensors/numpy.py,sha256=8ci56gDXetlYHH1-Nru83auiUVi-Q1P9bKvfsdkLKPw,5215
+safetensors/paddle.py,sha256=EhXpflqrhKr_NFh4jxV9SUnW0B1vcX_KdPdTqcytrDs,9011
+safetensors/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+safetensors/tensorflow.py,sha256=DajI3qkz00Zy2h7jublSAvTaD51QOPdaIgKQIeSiCRs,4042
+safetensors/torch.py,sha256=CLVWgWQdLm_tVzhRPaeihBHt-4iGAtUW5fY2ys3TyMc,19160

python_env/lib/site-packages/safetensors-0.7.0.dist-info/WHEEL

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

python_env/lib/site-packages/safetensors-0.7.0.dist-info/licenses/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

python_env/lib/site-packages/safetensors/__init__.py

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

python_env/lib/site-packages/safetensors/__init__.pyi

@@ -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.
+    """

python_env/lib/site-packages/safetensors/_safetensors_rust.pyd

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f203cb23bd0e572173eb37bc53faa5487d500f9bc2540901d0da466c5abadaf0
+size 738304

python_env/lib/site-packages/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.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

python_env/lib/site-packages/safetensors/mlx.py

@@ -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

python_env/lib/site-packages/safetensors/numpy.py

@@ -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}")

python_env/lib/site-packages/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()
+    }

python_env/lib/site-packages/safetensors/tensorflow.py

@@ -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

python_env/lib/site-packages/safetensors/torch.py

@@ -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()
+    }