safetensors 0.4.4rc0__tar.gz → 0.4.6.dev0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of safetensors might be problematic. Click here for more details.

Files changed (54) hide show
  1. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/PKG-INFO +2 -2
  2. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/Cargo.lock +45 -38
  3. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/Cargo.toml +2 -2
  4. safetensors-0.4.6.dev0/bindings/python/README.md +60 -0
  5. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/convert_all.py +1 -1
  6. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0/bindings/python}/py_src/safetensors/__init__.pyi +1 -1
  7. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/src/lib.rs +22 -6
  8. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/tests/test_simple.py +20 -0
  9. {safetensors-0.4.4rc0/bindings/python → safetensors-0.4.6.dev0}/py_src/safetensors/__init__.pyi +1 -1
  10. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/pyproject.toml +1 -1
  11. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/safetensors/Cargo.toml +2 -2
  12. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/safetensors/src/tensor.rs +36 -4
  13. safetensors-0.4.4rc0/safetensors/README.md +0 -193
  14. {safetensors-0.4.4rc0/bindings/python → safetensors-0.4.6.dev0}/README.md +0 -0
  15. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/.gitignore +0 -0
  16. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/MANIFEST.in +0 -0
  17. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/Makefile +0 -0
  18. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/benches/test_flax.py +0 -0
  19. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/benches/test_mlx.py +0 -0
  20. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/benches/test_paddle.py +0 -0
  21. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/benches/test_pt.py +0 -0
  22. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/benches/test_tf.py +0 -0
  23. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/convert.py +0 -0
  24. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/fuzz.py +0 -0
  25. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/py_src/safetensors/__init__.py +0 -0
  26. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/py_src/safetensors/flax.py +0 -0
  27. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/py_src/safetensors/mlx.py +0 -0
  28. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/py_src/safetensors/numpy.py +0 -0
  29. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/py_src/safetensors/paddle.py +0 -0
  30. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/py_src/safetensors/py.typed +0 -0
  31. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/py_src/safetensors/tensorflow.py +0 -0
  32. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/py_src/safetensors/torch.py +0 -0
  33. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/setup.cfg +0 -0
  34. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/stub.py +0 -0
  35. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/tests/data/__init__.py +0 -0
  36. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/tests/test_flax_comparison.py +0 -0
  37. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/tests/test_mlx_comparison.py +0 -0
  38. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/tests/test_paddle_comparison.py +0 -0
  39. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/tests/test_pt_comparison.py +0 -0
  40. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/tests/test_pt_model.py +0 -0
  41. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/bindings/python/tests/test_tf_comparison.py +0 -0
  42. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/py_src/safetensors/__init__.py +0 -0
  43. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/py_src/safetensors/flax.py +0 -0
  44. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/py_src/safetensors/mlx.py +0 -0
  45. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/py_src/safetensors/numpy.py +0 -0
  46. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/py_src/safetensors/paddle.py +0 -0
  47. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/py_src/safetensors/py.typed +0 -0
  48. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/py_src/safetensors/tensorflow.py +0 -0
  49. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/py_src/safetensors/torch.py +0 -0
  50. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/safetensors/LICENSE +0 -0
  51. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0/safetensors}/README.md +0 -0
  52. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/safetensors/benches/benchmark.rs +0 -0
  53. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/safetensors/src/lib.rs +0 -0
  54. {safetensors-0.4.4rc0 → safetensors-0.4.6.dev0}/safetensors/src/slice.rs +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: safetensors
3
- Version: 0.4.4rc0
3
+ Version: 0.4.6.dev0
4
4
  Classifier: Development Status :: 5 - Production/Stable
5
5
  Classifier: Intended Audience :: Developers
6
6
  Classifier: Intended Audience :: Education
@@ -20,7 +20,7 @@ Requires-Dist: torch >=1.10 ; extra == 'torch'
20
20
  Requires-Dist: safetensors[numpy] ; extra == 'tensorflow'
21
21
  Requires-Dist: tensorflow >=2.11.0 ; extra == 'tensorflow'
22
22
  Requires-Dist: safetensors[numpy] ; extra == 'pinned-tf'
23
- Requires-Dist: tensorflow ==2.11.0 ; extra == 'pinned-tf'
23
+ Requires-Dist: tensorflow ==2.18.0 ; extra == 'pinned-tf'
24
24
  Requires-Dist: safetensors[numpy] ; extra == 'jax'
25
25
  Requires-Dist: flax >=0.6.3 ; extra == 'jax'
26
26
  Requires-Dist: jax >=0.3.25 ; extra == 'jax'
@@ -4,9 +4,9 @@ version = 3
4
4
 
5
5
  [[package]]
6
6
  name = "autocfg"
7
- version = "1.3.0"
7
+ version = "1.4.0"
8
8
  source = "registry+https://github.com/rust-lang/crates.io-index"
9
- checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0"
9
+ checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26"
10
10
 
11
11
  [[package]]
12
12
  name = "cfg-if"
@@ -34,15 +34,21 @@ checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b"
34
34
 
35
35
  [[package]]
36
36
  name = "libc"
37
- version = "0.2.155"
37
+ version = "0.2.161"
38
38
  source = "registry+https://github.com/rust-lang/crates.io-index"
39
- checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c"
39
+ checksum = "8e9489c2807c139ffd9c1794f4af0ebe86a828db53ecdc7fea2111d0fed085d1"
40
+
41
+ [[package]]
42
+ name = "memchr"
43
+ version = "2.7.4"
44
+ source = "registry+https://github.com/rust-lang/crates.io-index"
45
+ checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
40
46
 
41
47
  [[package]]
42
48
  name = "memmap2"
43
- version = "0.9.4"
49
+ version = "0.9.5"
44
50
  source = "registry+https://github.com/rust-lang/crates.io-index"
45
- checksum = "fe751422e4a8caa417e13c3ea66452215d7d63e19e604f4980461212f3ae1322"
51
+ checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f"
46
52
  dependencies = [
47
53
  "libc",
48
54
  ]
@@ -58,30 +64,30 @@ dependencies = [
58
64
 
59
65
  [[package]]
60
66
  name = "once_cell"
61
- version = "1.19.0"
67
+ version = "1.20.2"
62
68
  source = "registry+https://github.com/rust-lang/crates.io-index"
63
- checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92"
69
+ checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775"
64
70
 
65
71
  [[package]]
66
72
  name = "portable-atomic"
67
- version = "1.6.0"
73
+ version = "1.9.0"
68
74
  source = "registry+https://github.com/rust-lang/crates.io-index"
69
- checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0"
75
+ checksum = "cc9c68a3f6da06753e9335d63e27f6b9754dd1920d941135b7ea8224f141adb2"
70
76
 
71
77
  [[package]]
72
78
  name = "proc-macro2"
73
- version = "1.0.85"
79
+ version = "1.0.89"
74
80
  source = "registry+https://github.com/rust-lang/crates.io-index"
75
- checksum = "22244ce15aa966053a896d1accb3a6e68469b97c7f33f284b99f0d576879fc23"
81
+ checksum = "f139b0662de085916d1fb67d2b4169d1addddda1919e696f3252b740b629986e"
76
82
  dependencies = [
77
83
  "unicode-ident",
78
84
  ]
79
85
 
80
86
  [[package]]
81
87
  name = "pyo3"
82
- version = "0.22.2"
88
+ version = "0.22.6"
83
89
  source = "registry+https://github.com/rust-lang/crates.io-index"
84
- checksum = "831e8e819a138c36e212f3af3fd9eeffed6bf1510a805af35b0edee5ffa59433"
90
+ checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884"
85
91
  dependencies = [
86
92
  "cfg-if",
87
93
  "indoc",
@@ -97,9 +103,9 @@ dependencies = [
97
103
 
98
104
  [[package]]
99
105
  name = "pyo3-build-config"
100
- version = "0.22.2"
106
+ version = "0.22.6"
101
107
  source = "registry+https://github.com/rust-lang/crates.io-index"
102
- checksum = "1e8730e591b14492a8945cdff32f089250b05f5accecf74aeddf9e8272ce1fa8"
108
+ checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38"
103
109
  dependencies = [
104
110
  "once_cell",
105
111
  "target-lexicon",
@@ -107,9 +113,9 @@ dependencies = [
107
113
 
108
114
  [[package]]
109
115
  name = "pyo3-ffi"
110
- version = "0.22.2"
116
+ version = "0.22.6"
111
117
  source = "registry+https://github.com/rust-lang/crates.io-index"
112
- checksum = "5e97e919d2df92eb88ca80a037969f44e5e70356559654962cbb3316d00300c6"
118
+ checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636"
113
119
  dependencies = [
114
120
  "libc",
115
121
  "pyo3-build-config",
@@ -117,9 +123,9 @@ dependencies = [
117
123
 
118
124
  [[package]]
119
125
  name = "pyo3-macros"
120
- version = "0.22.2"
126
+ version = "0.22.6"
121
127
  source = "registry+https://github.com/rust-lang/crates.io-index"
122
- checksum = "eb57983022ad41f9e683a599f2fd13c3664d7063a3ac5714cae4b7bee7d3f206"
128
+ checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453"
123
129
  dependencies = [
124
130
  "proc-macro2",
125
131
  "pyo3-macros-backend",
@@ -129,9 +135,9 @@ dependencies = [
129
135
 
130
136
  [[package]]
131
137
  name = "pyo3-macros-backend"
132
- version = "0.22.2"
138
+ version = "0.22.6"
133
139
  source = "registry+https://github.com/rust-lang/crates.io-index"
134
- checksum = "ec480c0c51ddec81019531705acac51bcdbeae563557c982aa8263bb96880372"
140
+ checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe"
135
141
  dependencies = [
136
142
  "heck",
137
143
  "proc-macro2",
@@ -142,9 +148,9 @@ dependencies = [
142
148
 
143
149
  [[package]]
144
150
  name = "quote"
145
- version = "1.0.36"
151
+ version = "1.0.37"
146
152
  source = "registry+https://github.com/rust-lang/crates.io-index"
147
- checksum = "0fa76aaf39101c457836aec0ce2316dbdc3ab723cdda1c6bd4e6ad4208acaca7"
153
+ checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af"
148
154
  dependencies = [
149
155
  "proc-macro2",
150
156
  ]
@@ -157,7 +163,7 @@ checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f"
157
163
 
158
164
  [[package]]
159
165
  name = "safetensors"
160
- version = "0.4.4-rc.0"
166
+ version = "0.4.6-dev.0"
161
167
  dependencies = [
162
168
  "serde",
163
169
  "serde_json",
@@ -165,7 +171,7 @@ dependencies = [
165
171
 
166
172
  [[package]]
167
173
  name = "safetensors-python"
168
- version = "0.4.4-rc.0"
174
+ version = "0.4.6-dev.0"
169
175
  dependencies = [
170
176
  "memmap2",
171
177
  "pyo3",
@@ -175,18 +181,18 @@ dependencies = [
175
181
 
176
182
  [[package]]
177
183
  name = "serde"
178
- version = "1.0.203"
184
+ version = "1.0.214"
179
185
  source = "registry+https://github.com/rust-lang/crates.io-index"
180
- checksum = "7253ab4de971e72fb7be983802300c30b5a7f0c2e56fab8abfc6a214307c0094"
186
+ checksum = "f55c3193aca71c12ad7890f1785d2b73e1b9f63a0bbc353c08ef26fe03fc56b5"
181
187
  dependencies = [
182
188
  "serde_derive",
183
189
  ]
184
190
 
185
191
  [[package]]
186
192
  name = "serde_derive"
187
- version = "1.0.203"
193
+ version = "1.0.214"
188
194
  source = "registry+https://github.com/rust-lang/crates.io-index"
189
- checksum = "500cbc0ebeb6f46627f50f3f5811ccf6bf00643be300b4c3eabc0ef55dc5b5ba"
195
+ checksum = "de523f781f095e28fa605cdce0f8307e451cc0fd14e2eb4cd2e98a355b147766"
190
196
  dependencies = [
191
197
  "proc-macro2",
192
198
  "quote",
@@ -195,20 +201,21 @@ dependencies = [
195
201
 
196
202
  [[package]]
197
203
  name = "serde_json"
198
- version = "1.0.117"
204
+ version = "1.0.132"
199
205
  source = "registry+https://github.com/rust-lang/crates.io-index"
200
- checksum = "455182ea6142b14f93f4bc5320a2b31c1f266b66a4a5c858b013302a5d8cbfc3"
206
+ checksum = "d726bfaff4b320266d395898905d0eba0345aae23b54aee3a737e260fd46db03"
201
207
  dependencies = [
202
208
  "itoa",
209
+ "memchr",
203
210
  "ryu",
204
211
  "serde",
205
212
  ]
206
213
 
207
214
  [[package]]
208
215
  name = "syn"
209
- version = "2.0.66"
216
+ version = "2.0.87"
210
217
  source = "registry+https://github.com/rust-lang/crates.io-index"
211
- checksum = "c42f3f41a2de00b01c0aaad383c5a45241efc8b2d1eda5661812fda5f3cdcff5"
218
+ checksum = "25aa4ce346d03a6dcd68dd8b4010bcb74e54e62c90c573f394c46eae99aba32d"
212
219
  dependencies = [
213
220
  "proc-macro2",
214
221
  "quote",
@@ -217,15 +224,15 @@ dependencies = [
217
224
 
218
225
  [[package]]
219
226
  name = "target-lexicon"
220
- version = "0.12.14"
227
+ version = "0.12.16"
221
228
  source = "registry+https://github.com/rust-lang/crates.io-index"
222
- checksum = "e1fc403891a21bcfb7c37834ba66a547a8f402146eba7265b5a6d88059c9ff2f"
229
+ checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
223
230
 
224
231
  [[package]]
225
232
  name = "unicode-ident"
226
- version = "1.0.12"
233
+ version = "1.0.13"
227
234
  source = "registry+https://github.com/rust-lang/crates.io-index"
228
- checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
235
+ checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe"
229
236
 
230
237
  [[package]]
231
238
  name = "unindent"
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "safetensors-python"
3
- version = "0.4.4-rc.0"
3
+ version = "0.4.6-dev.0"
4
4
  edition = "2021"
5
5
 
6
6
  # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -9,7 +9,7 @@ name = "safetensors_rust"
9
9
  crate-type = ["cdylib"]
10
10
 
11
11
  [dependencies]
12
- pyo3 = { version = "0.22" }
12
+ pyo3 = { version = "0.22", features = ["abi3", "abi3-py38"] }
13
13
  memmap2 = "0.9"
14
14
  serde_json = "1.0"
15
15
 
@@ -0,0 +1,60 @@
1
+ ## Installation
2
+
3
+ ```
4
+ pip install safetensors
5
+ ```
6
+
7
+
8
+ ## Usage
9
+
10
+ ### Numpy
11
+
12
+ ```python
13
+ from safetensors.numpy import save_file, load_file
14
+ import numpy as np
15
+
16
+ tensors = {
17
+ "a": np.zeros((2, 2)),
18
+ "b": np.zeros((2, 3), dtype=np.uint8)
19
+ }
20
+
21
+ save_file(tensors, "./model.safetensors")
22
+
23
+
24
+ # Now loading
25
+ loaded = load_file("./model.safetensors")
26
+ ```
27
+
28
+ ### Torch
29
+
30
+ ```python
31
+ from safetensors.torch import save_file, load_file
32
+ import torch
33
+
34
+ tensors = {
35
+ "a": torch.zeros((2, 2)),
36
+ "b": torch.zeros((2, 3), dtype=torch.uint8)
37
+ }
38
+
39
+ save_file(tensors, "./model.safetensors")
40
+
41
+
42
+ # Now loading
43
+ loaded = load_file("./model.safetensors")
44
+ ```
45
+
46
+ ### Developing
47
+
48
+ ```
49
+ # inside ./safetensors/bindings/python
50
+ pip install .[dev]
51
+ ```
52
+ Should be enough to install this library locally.
53
+
54
+ ### Testing
55
+
56
+ ```
57
+ # inside ./safetensors/bindings/python
58
+ pip install .[dev]
59
+ pytest -sv tests/
60
+ ```
@@ -16,7 +16,7 @@ if __name__ == "__main__":
16
16
  correct = 0
17
17
  errors = set()
18
18
  for model in models:
19
- model = api.model_info(model.modelId, files_metadata=True)
19
+ model = api.model_info(model.id, files_metadata=True)
20
20
  size = None
21
21
  for sibling in model.siblings:
22
22
  if sibling.rfilename == "pytorch_model.bin":
@@ -62,7 +62,7 @@ class safe_open:
62
62
  The filename to open
63
63
 
64
64
  framework (:obj:`str`):
65
- The framework you want you tensors in. Supported values:
65
+ The framework you want your tensors in. Supported values:
66
66
  `pt`, `tf`, `flax`, `numpy`.
67
67
 
68
68
  device (:obj:`str`, defaults to :obj:`"cpu"`):
@@ -5,8 +5,7 @@ use pyo3::exceptions::{PyException, PyFileNotFoundError};
5
5
  use pyo3::prelude::*;
6
6
  use pyo3::sync::GILOnceCell;
7
7
  use pyo3::types::IntoPyDict;
8
- use pyo3::types::PySlice;
9
- use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList};
8
+ use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList, PySlice};
10
9
  use pyo3::Bound as PyBound;
11
10
  use pyo3::{intern, PyErr};
12
11
  use safetensors::slice::TensorIndexer;
@@ -65,8 +64,8 @@ fn prepare(tensor_dict: HashMap<String, PyBound<PyDict>>) -> PyResult<HashMap<St
65
64
  let pydtype = tensor_desc.get_item("dtype")?.ok_or_else(|| {
66
65
  SafetensorError::new_err(format!("Missing `dtype` in {tensor_desc:?}"))
67
66
  })?;
68
- let dtype: &str = pydtype.extract()?;
69
- let dtype = match dtype {
67
+ let dtype: String = pydtype.extract()?;
68
+ let dtype = match dtype.as_ref() {
70
69
  "bool" => Dtype::BOOL,
71
70
  "int8" => Dtype::I8,
72
71
  "uint8" => Dtype::U8,
@@ -842,10 +841,27 @@ impl PySafeSlice {
842
841
  pub fn __getitem__(&self, slices: &PyBound<'_, PyAny>) -> PyResult<PyObject> {
843
842
  match &self.storage.as_ref() {
844
843
  Storage::Mmap(mmap) => {
845
- let slices: Slice = slices.extract()?;
844
+ let pyslices = slices;
845
+ let slices: Slice = pyslices.extract()?;
846
+ let is_list = pyslices.is_instance_of::<PyList>();
846
847
  let slices: Vec<SliceIndex> = match slices {
847
848
  Slice::Slice(slice) => vec![slice],
848
- Slice::Slices(slices) => slices,
849
+ Slice::Slices(slices) => {
850
+ if slices.is_empty() && is_list {
851
+ vec![SliceIndex::Slice(PySlice::new_bound(
852
+ pyslices.py(),
853
+ 0,
854
+ 0,
855
+ 0,
856
+ ))]
857
+ } else if is_list {
858
+ return Err(SafetensorError::new_err(
859
+ "Non empty lists are not implemented",
860
+ ));
861
+ } else {
862
+ slices
863
+ }
864
+ }
849
865
  };
850
866
  let data = &mmap[self.info.data_offsets.0 + self.offset
851
867
  ..self.info.data_offsets.1 + self.offset];
@@ -246,6 +246,10 @@ class ReadmeTestCase(unittest.TestCase):
246
246
  self.assertEqual(list(tensor.shape), [10, 5])
247
247
  torch.testing.assert_close(tensor, A)
248
248
 
249
+ tensor = slice_[tuple()]
250
+ self.assertEqual(list(tensor.shape), [10, 5])
251
+ torch.testing.assert_close(tensor, A)
252
+
249
253
  tensor = slice_[:2]
250
254
  self.assertEqual(list(tensor.shape), [2, 5])
251
255
  torch.testing.assert_close(tensor, A[:2])
@@ -270,6 +274,10 @@ class ReadmeTestCase(unittest.TestCase):
270
274
  self.assertEqual(list(tensor.shape), [8])
271
275
  torch.testing.assert_close(tensor, A[2:, -1])
272
276
 
277
+ tensor = slice_[list()]
278
+ self.assertEqual(list(tensor.shape), [0, 5])
279
+ torch.testing.assert_close(tensor, A[list()])
280
+
273
281
  def test_numpy_slice(self):
274
282
  A = np.random.rand(10, 5)
275
283
  tensors = {
@@ -284,6 +292,10 @@ class ReadmeTestCase(unittest.TestCase):
284
292
  self.assertEqual(list(tensor.shape), [10, 5])
285
293
  self.assertTrue(np.allclose(tensor, A))
286
294
 
295
+ tensor = slice_[tuple()]
296
+ self.assertEqual(list(tensor.shape), [10, 5])
297
+ self.assertTrue(np.allclose(tensor, A))
298
+
287
299
  tensor = slice_[:2]
288
300
  self.assertEqual(list(tensor.shape), [2, 5])
289
301
  self.assertTrue(np.allclose(tensor, A[:2]))
@@ -312,10 +324,18 @@ class ReadmeTestCase(unittest.TestCase):
312
324
  self.assertEqual(list(tensor.shape), [8])
313
325
  self.assertTrue(np.allclose(tensor, A[2:, -5]))
314
326
 
327
+ tensor = slice_[list()]
328
+ self.assertEqual(list(tensor.shape), [0, 5])
329
+ self.assertTrue(np.allclose(tensor, A[list()]))
330
+
315
331
  with self.assertRaises(SafetensorError) as cm:
316
332
  tensor = slice_[2:, -6]
317
333
  self.assertEqual(str(cm.exception), "Invalid index -6 for dimension 1 of size 5")
318
334
 
335
+ with self.assertRaises(SafetensorError) as cm:
336
+ tensor = slice_[[0, 1]]
337
+ self.assertEqual(str(cm.exception), "Non empty lists are not implemented")
338
+
319
339
  with self.assertRaises(SafetensorError) as cm:
320
340
  tensor = slice_[2:, 20]
321
341
  self.assertEqual(
@@ -62,7 +62,7 @@ class safe_open:
62
62
  The filename to open
63
63
 
64
64
  framework (:obj:`str`):
65
- The framework you want you tensors in. Supported values:
65
+ The framework you want your tensors in. Supported values:
66
66
  `pt`, `tf`, `flax`, `numpy`.
67
67
 
68
68
  device (:obj:`str`, defaults to :obj:`"cpu"`):
@@ -42,7 +42,7 @@ tensorflow = [
42
42
  # pinning tf version 2.11.0 for doc-builder
43
43
  pinned-tf = [
44
44
  "safetensors[numpy]",
45
- "tensorflow==2.11.0",
45
+ "tensorflow==2.18.0",
46
46
  ]
47
47
  jax = [
48
48
  "safetensors[numpy]",
@@ -1,13 +1,13 @@
1
1
  [package]
2
2
  name = "safetensors"
3
- version = "0.4.4-rc.0"
3
+ version = "0.4.6-dev.0"
4
4
  edition = "2021"
5
5
  homepage = "https://github.com/huggingface/safetensors"
6
6
  repository = "https://github.com/huggingface/safetensors"
7
7
  documentation = "https://docs.rs/safetensors/"
8
8
  license = "Apache-2.0"
9
9
  keywords = ["safetensors", "huggingface", "Tensors", "Pytorch", "Tensorflow"]
10
- readme = "./README.md"
10
+ readme = "README.md"
11
11
  description = """
12
12
  Provides functions to read and write safetensors which aim to be safer than
13
13
  their PyTorch counterpart.
@@ -345,10 +345,10 @@ impl<'data> SafeTensors<'data> {
345
345
  Ok(Self { metadata, data })
346
346
  }
347
347
 
348
- /// Allow the user to iterate over tensors within the SafeTensors.
348
+ /// Returns the tensors contained within the SafeTensors.
349
349
  /// The tensors returned are merely views and the data is not owned by this
350
350
  /// structure.
351
- pub fn tensors(&self) -> Vec<(String, TensorView<'_>)> {
351
+ pub fn tensors(&self) -> Vec<(String, TensorView<'data>)> {
352
352
  let mut tensors = Vec::with_capacity(self.metadata.index_map.len());
353
353
  for (name, &index) in &self.metadata.index_map {
354
354
  let info = &self.metadata.tensors[index];
@@ -362,10 +362,27 @@ impl<'data> SafeTensors<'data> {
362
362
  tensors
363
363
  }
364
364
 
365
+ /// Returns an iterator over the tensors contained within the SafeTensors.
366
+ /// The tensors returned are merely views and the data is not owned by this
367
+ /// structure.
368
+ pub fn iter<'a>(&'a self) -> impl Iterator<Item = (&'a str, TensorView<'data>)> {
369
+ self.metadata.index_map.iter().map(|(name, &idx)| {
370
+ let info = &self.metadata.tensors[idx];
371
+ (
372
+ name.as_str(),
373
+ TensorView {
374
+ dtype: info.dtype,
375
+ shape: info.shape.clone(),
376
+ data: &self.data[info.data_offsets.0..info.data_offsets.1],
377
+ },
378
+ )
379
+ })
380
+ }
381
+
365
382
  /// Allow the user to get a specific tensor within the SafeTensors.
366
383
  /// The tensor returned is merely a view and the data is not owned by this
367
384
  /// structure.
368
- pub fn tensor(&self, tensor_name: &str) -> Result<TensorView<'_>, SafeTensorError> {
385
+ pub fn tensor(&self, tensor_name: &str) -> Result<TensorView<'data>, SafeTensorError> {
369
386
  if let Some(index) = &self.metadata.index_map.get(tensor_name) {
370
387
  if let Some(info) = &self.metadata.tensors.get(**index) {
371
388
  Ok(TensorView {
@@ -541,7 +558,7 @@ impl Metadata {
541
558
  /// A view of a Tensor within the file.
542
559
  /// Contains references to data within the full byte-buffer
543
560
  /// And is thus a readable view of a single tensor
544
- #[derive(Debug, PartialEq, Eq)]
561
+ #[derive(Debug, PartialEq, Eq, Clone)]
545
562
  pub struct TensorView<'data> {
546
563
  dtype: Dtype,
547
564
  shape: Vec<usize>,
@@ -1038,6 +1055,21 @@ mod tests {
1038
1055
  assert_eq!(tensor.data(), b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0");
1039
1056
  }
1040
1057
 
1058
+ #[test]
1059
+ fn test_lifetimes() {
1060
+ let serialized = b"<\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,2],\"data_offsets\":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00";
1061
+
1062
+ let tensor = {
1063
+ let loaded = SafeTensors::deserialize(serialized).unwrap();
1064
+ loaded.tensor("test").unwrap()
1065
+ };
1066
+
1067
+ assert_eq!(tensor.shape(), vec![2, 2]);
1068
+ assert_eq!(tensor.dtype(), Dtype::I32);
1069
+ // 16 bytes
1070
+ assert_eq!(tensor.data(), b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0");
1071
+ }
1072
+
1041
1073
  #[test]
1042
1074
  fn test_json_attack() {
1043
1075
  let mut tensors = HashMap::new();
@@ -1,193 +0,0 @@
1
- <p align="center">
2
- <picture>
3
- <source media="(prefers-color-scheme: dark)" srcset="https://huggingface.co/datasets/safetensors/assets/raw/main/banner-dark.svg">
4
- <source media="(prefers-color-scheme: light)" srcset="https://huggingface.co/datasets/safetensors/assets/raw/main/banner-light.svg">
5
- <img alt="Hugging Face Safetensors Library" src="https://huggingface.co/datasets/safetensors/assets/raw/main/banner-light.svg" style="max-width: 100%;">
6
- </picture>
7
- <br/>
8
- <br/>
9
- </p>
10
-
11
- Python
12
- [![Pypi](https://img.shields.io/pypi/v/safetensors.svg)](https://pypi.org/pypi/safetensors/)
13
- [![Documentation](https://img.shields.io/website/http/huggingface.co/docs/safetensors/index.svg?label=docs)](https://huggingface.co/docs/safetensors/index)
14
- [![Codecov](https://codecov.io/github/huggingface/safetensors/coverage.svg?branch=main)](https://codecov.io/gh/huggingface/safetensors)
15
- [![Downloads](https://static.pepy.tech/badge/safetensors/month)](https://pepy.tech/project/safetensors)
16
-
17
- Rust
18
- [![Crates.io](https://img.shields.io/crates/v/safetensors.svg)](https://crates.io/crates/safetensors)
19
- [![Documentation](https://docs.rs/safetensors/badge.svg)](https://docs.rs/safetensors/)
20
- [![Codecov](https://codecov.io/github/huggingface/safetensors/coverage.svg?branch=main)](https://codecov.io/gh/huggingface/safetensors)
21
- [![Dependency status](https://deps.rs/repo/github/huggingface/safetensors/status.svg?path=safetensors)](https://deps.rs/repo/github/huggingface/safetensors?path=safetensors)
22
-
23
- # safetensors
24
-
25
- ## Safetensors
26
-
27
- This repository implements a new simple format for storing tensors
28
- safely (as opposed to pickle) and that is still fast (zero-copy).
29
-
30
- ### Installation
31
- #### Pip
32
-
33
- You can install safetensors via the pip manager:
34
-
35
- ```bash
36
- pip install safetensors
37
- ```
38
-
39
- #### From source
40
-
41
- For the sources, you need Rust
42
-
43
- ```bash
44
- # Install Rust
45
- curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
46
- # Make sure it's up to date and using stable channel
47
- rustup update
48
- git clone https://github.com/huggingface/safetensors
49
- cd safetensors/bindings/python
50
- pip install setuptools_rust
51
- pip install -e .
52
- ```
53
-
54
- ### Getting started
55
-
56
- ```python
57
- import torch
58
- from safetensors import safe_open
59
- from safetensors.torch import save_file
60
-
61
- tensors = {
62
- "weight1": torch.zeros((1024, 1024)),
63
- "weight2": torch.zeros((1024, 1024))
64
- }
65
- save_file(tensors, "model.safetensors")
66
-
67
- tensors = {}
68
- with safe_open("model.safetensors", framework="pt", device="cpu") as f:
69
- for key in f.keys():
70
- tensors[key] = f.get_tensor(key)
71
- ```
72
-
73
- [Python documentation](https://huggingface.co/docs/safetensors/index)
74
-
75
-
76
- ### Format
77
-
78
- - 8 bytes: `N`, an unsigned little-endian 64-bit integer, containing the size of the header
79
- - N bytes: a JSON UTF-8 string representing the header.
80
- - The header data MUST begin with a `{` character (0x7B).
81
- - The header data MAY be trailing padded with whitespace (0x20).
82
- - The header is a dict like `{"TENSOR_NAME": {"dtype": "F16", "shape": [1, 16, 256], "data_offsets": [BEGIN, END]}, "NEXT_TENSOR_NAME": {...}, ...}`,
83
- - `data_offsets` point to the tensor data relative to the beginning of the byte buffer (i.e. not an absolute position in the file),
84
- with `BEGIN` as the starting offset and `END` as the one-past offset (so total tensor byte size = `END - BEGIN`).
85
- - A special key `__metadata__` is allowed to contain free form string-to-string map. Arbitrary JSON is not allowed, all values must be strings.
86
- - Rest of the file: byte-buffer.
87
-
88
- Notes:
89
- - Duplicate keys are disallowed. Not all parsers may respect this.
90
- - In general the subset of JSON is implicitly decided by `serde_json` for
91
- this library. Anything obscure might be modified at a later time, that odd ways
92
- to represent integer, newlines and escapes in utf-8 strings. This would only
93
- be done for safety concerns
94
- - Tensor values are not checked against, in particular NaN and +/-Inf could
95
- be in the file
96
- - Empty tensors (tensors with 1 dimension being 0) are allowed.
97
- They are not storing any data in the databuffer, yet retaining size in the header.
98
- They don't really bring a lot of values but are accepted since they are valid tensors
99
- from traditional tensor libraries perspective (torch, tensorflow, numpy, ..).
100
- - 0-rank Tensors (tensors with shape `[]`) are allowed, they are merely a scalar.
101
- - The byte buffer needs to be entirely indexed, and cannot contain holes. This prevents
102
- the creation of polyglot files.
103
- - Endianness: Little-endian.
104
- moment.
105
- - Order: 'C' or row-major.
106
-
107
-
108
- ### Yet another format ?
109
-
110
- The main rationale for this crate is to remove the need to use
111
- `pickle` on `PyTorch` which is used by default.
112
- There are other formats out there used by machine learning and more general
113
- formats.
114
-
115
-
116
- Let's take a look at alternatives and why this format is deemed interesting.
117
- This is my very personal and probably biased view:
118
-
119
- | Format | Safe | Zero-copy | Lazy loading | No file size limit | Layout control | Flexibility | Bfloat16/Fp8
120
- | ----------------------- | --- | --- | --- | --- | --- | --- | --- |
121
- | pickle (PyTorch) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 |
122
- | H5 (Tensorflow) | 🗸 | ✗ | 🗸 | 🗸 | ~ | ~ | ✗ |
123
- | SavedModel (Tensorflow) | 🗸 | ✗ | ✗ | 🗸 | 🗸 | ✗ | 🗸 |
124
- | MsgPack (flax) | 🗸 | 🗸 | ✗ | 🗸 | ✗ | ✗ | 🗸 |
125
- | Protobuf (ONNX) | 🗸 | ✗ | ✗ | ✗ | ✗ | ✗ | 🗸 |
126
- | Cap'n'Proto | 🗸 | 🗸 | ~ | 🗸 | 🗸 | ~ | ✗ |
127
- | Arrow | ? | ? | ? | ? | ? | ? | ✗ |
128
- | Numpy (npy,npz) | 🗸 | ? | ? | ✗ | 🗸 | ✗ | ✗ |
129
- | pdparams (Paddle) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 |
130
- | SafeTensors | 🗸 | 🗸 | 🗸 | 🗸 | 🗸 | ✗ | 🗸 |
131
-
132
- - Safe: Can I use a file randomly downloaded and expect not to run arbitrary code ?
133
- - Zero-copy: Does reading the file require more memory than the original file ?
134
- - Lazy loading: Can I inspect the file without loading everything ? And loading only
135
- some tensors in it without scanning the whole file (distributed setting) ?
136
- - Layout control: Lazy loading, is not necessarily enough since if the information about tensors is spread out in your file, then even if the information is lazily accessible you might have to access most of your file to read the available tensors (incurring many DISK -> RAM copies). Controlling the layout to keep fast access to single tensors is important.
137
- - No file size limit: Is there a limit to the file size ?
138
- - Flexibility: Can I save custom code in the format and be able to use it later with zero extra code ? (~ means we can store more than pure tensors, but no custom code)
139
- - Bfloat16/Fp8: Does the format support native bfloat16/fp8 (meaning no weird workarounds are
140
- necessary)? This is becoming increasingly important in the ML world.
141
-
142
-
143
- ### Main oppositions
144
-
145
- - Pickle: Unsafe, runs arbitrary code
146
- - H5: Apparently now discouraged for TF/Keras. Seems like a great fit otherwise actually. Some classic use after free issues: <https://www.cvedetails.com/vulnerability-list/vendor_id-15991/product_id-35054/Hdfgroup-Hdf5.html>. On a very different level than pickle security-wise. Also 210k lines of code vs ~400 lines for this lib currently.
147
- - SavedModel: Tensorflow specific (it contains TF graph information).
148
- - MsgPack: No layout control to enable lazy loading (important for loading specific parts in distributed setting)
149
- - Protobuf: Hard 2Go max file size limit
150
- - Cap'n'proto: Float16 support is not present [link](https://capnproto.org/language.html#built-in-types) so using a manual wrapper over a byte-buffer would be necessary. Layout control seems possible but not trivial as buffers have limitations [link](https://stackoverflow.com/questions/48458839/capnproto-maximum-filesize).
151
- - Numpy (npz): No `bfloat16` support. Vulnerable to zip bombs (DOS). Not zero-copy.
152
- - Arrow: No `bfloat16` support.
153
-
154
- ### Notes
155
-
156
- - Zero-copy: No format is really zero-copy in ML, it needs to go from disk to RAM/GPU RAM (that takes time). On CPU, if the file is already in cache, then it can
157
- truly be zero-copy, whereas on GPU there is not such disk cache, so a copy is always required
158
- but you can bypass allocating all the tensors on CPU at any given point.
159
- SafeTensors is not zero-copy for the header. The choice of JSON is pretty arbitrary, but since deserialization is <<< of the time required to load the actual tensor data and is readable I went that way, (also space is <<< to the tensor data).
160
-
161
- - Endianness: Little-endian. This can be modified later, but it feels really unnecessary at the
162
- moment.
163
- - Order: 'C' or row-major. This seems to have won. We can add that information later if needed.
164
- - Stride: No striding, all tensors need to be packed before being serialized. I have yet to see a case where it seems useful to have a strided tensor stored in serialized format.
165
-
166
- ### Benefits
167
-
168
- Since we can invent a new format we can propose additional benefits:
169
-
170
- - Prevent DOS attacks: We can craft the format in such a way that it's almost
171
- impossible to use malicious files to DOS attack a user. Currently, there's a limit
172
- on the size of the header of 100MB to prevent parsing extremely large JSON.
173
- Also when reading the file, there's a guarantee that addresses in the file
174
- do not overlap in any way, meaning when you're loading a file you should never
175
- exceed the size of the file in memory
176
-
177
- - Faster load: PyTorch seems to be the fastest file to load out in the major
178
- ML formats. However, it does seem to have an extra copy on CPU, which we
179
- can bypass in this lib by using `torch.UntypedStorage.from_file`.
180
- Currently, CPU loading times are extremely fast with this lib compared to pickle.
181
- GPU loading times are as fast or faster than PyTorch equivalent.
182
- Loading first on CPU with memmapping with torch, and then moving all tensors to GPU seems
183
- to be faster too somehow (similar behavior in torch pickle)
184
-
185
- - Lazy loading: in distributed (multi-node or multi-gpu) settings, it's nice to be able to
186
- load only part of the tensors on the various models. For
187
- [BLOOM](https://huggingface.co/bigscience/bloom) using this format enabled
188
- to load the model on 8 GPUs from 10mn with regular PyTorch weights down to 45s.
189
- This really speeds up feedbacks loops when developing on the model. For instance
190
- you don't have to have separate copies of the weights when changing the distribution
191
- strategy (for instance Pipeline Parallelism vs Tensor Parallelism).
192
-
193
- License: Apache-2.0