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