tjh.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 rustPlatform,
7
8 # optional-dependencies
9 numpy,
10 torch,
11 tensorflow,
12 flax,
13 jax,
14 mlx,
15 paddlepaddle,
16 h5py,
17 huggingface-hub,
18 setuptools-rust,
19 pytest,
20 pytest-benchmark,
21 hypothesis,
22
23 # tests
24 pytestCheckHook,
25}:
26
27buildPythonPackage rec {
28 pname = "safetensors";
29 version = "0.6.2";
30 pyproject = true;
31
32 src = fetchFromGitHub {
33 owner = "huggingface";
34 repo = "safetensors";
35 tag = "v${version}";
36 hash = "sha256-IyKk29jMAbYW+16mrpqQWjnsmNFEvUwkB048AAx/Cvw=";
37 };
38
39 sourceRoot = "${src.name}/bindings/python";
40
41 cargoDeps = rustPlatform.fetchCargoVendor {
42 inherit
43 pname
44 version
45 src
46 sourceRoot
47 ;
48 hash = "sha256-+92fCILZwk/TknGXgR9lRN55WnmkgUJfCszFthstzXs=";
49 };
50
51 nativeBuildInputs = [
52 rustPlatform.cargoSetupHook
53 rustPlatform.maturinBuildHook
54 ];
55
56 optional-dependencies = lib.fix (self: {
57 numpy = [ numpy ];
58 torch = self.numpy ++ [
59 torch
60 ];
61 tensorflow = self.numpy ++ [
62 tensorflow
63 ];
64 pinned-tf = self.numpy ++ [
65 tensorflow
66 ];
67 jax = self.numpy ++ [
68 flax
69 jax
70 ];
71 mlx = [
72 mlx
73 ];
74 paddlepaddle = self.numpy ++ [
75 paddlepaddle
76 ];
77 testing = self.numpy ++ [
78 h5py
79 huggingface-hub
80 setuptools-rust
81 pytest
82 pytest-benchmark
83 hypothesis
84 ];
85 all = self.torch ++ self.numpy ++ self.pinned-tf ++ self.jax ++ self.paddlepaddle ++ self.testing;
86 dev = self.all;
87 });
88
89 nativeCheckInputs = [
90 h5py
91 numpy
92 pytestCheckHook
93 torch
94 ];
95
96 enabledTestPaths = [ "tests" ];
97
98 disabledTests = [
99 # AttributeError: module 'torch' has no attribute 'float4_e2m1fn_x2'
100 "test_odd_dtype_fp4"
101
102 # AssertionError: 'No such file or directory: notafile' != 'No such file or directory: "notafile"'
103 "test_file_not_found"
104
105 # AssertionError:
106 # 'Erro[41 chars] 5]: index 20 out of bounds for tensor dimension #1 of size 5'
107 # != 'Erro[41 chars] 5]: SliceOutOfRange { dim_index: 1, asked: 20, dim_size: 5 }'
108 "test_numpy_slice"
109 ];
110
111 # don't require PaddlePaddle (not in Nixpkgs), Flax, or Tensorflow (onerous) to run tests:
112 disabledTestPaths = [
113 "tests/test_flax_comparison.py"
114 "tests/test_paddle_comparison.py"
115 "tests/test_tf_comparison.py"
116 ]
117 ++ lib.optionals stdenv.hostPlatform.isDarwin [
118 # don't require mlx (not in Nixpkgs) to run tests
119 "tests/test_mlx_comparison.py"
120 ];
121
122 pythonImportsCheck = [ "safetensors" ];
123
124 meta = {
125 homepage = "https://github.com/huggingface/safetensors";
126 description = "Fast (zero-copy) and safe (unlike pickle) format for storing tensors";
127 changelog = "https://github.com/huggingface/safetensors/releases/tag/v${version}";
128 license = lib.licenses.asl20;
129 maintainers = with lib.maintainers; [ bcdarwin ];
130 };
131}