tjh.dev
1{
2 stdenv,
3 lib,
4 buildPythonPackage,
5 cargo,
6 fetchFromGitHub,
7 h5py,
8 numpy,
9 pythonOlder,
10 pytestCheckHook,
11 rustc,
12 rustPlatform,
13 setuptools-rust,
14 torch,
15 libiconv,
16}:
17
18buildPythonPackage rec {
19 pname = "safetensors";
20 version = "0.4.3";
21 pyproject = true;
22
23 disabled = pythonOlder "3.7";
24
25 src = fetchFromGitHub {
26 owner = "huggingface";
27 repo = "safetensors";
28 rev = "refs/tags/v${version}";
29 hash = "sha256-Rc+o7epQJ8qEvdgbFnGvXxBr/U4eULZwkKNEaPlJkyU=";
30 };
31
32 cargoDeps = rustPlatform.fetchCargoTarball {
33 inherit src;
34 sourceRoot = "${src.name}/bindings/python";
35 hash = "sha256-tzNEUvWgolSwX0t/JLgYcTEIv3/FiKxoTJ4VjFQs8AY=";
36 };
37
38 sourceRoot = "${src.name}/bindings/python";
39
40 nativeBuildInputs = [
41 setuptools-rust
42 cargo
43 rustc
44 rustPlatform.cargoSetupHook
45 rustPlatform.maturinBuildHook
46 ];
47
48 buildInputs = lib.optionals stdenv.isDarwin [ libiconv ];
49
50 nativeCheckInputs = [
51 h5py
52 numpy
53 pytestCheckHook
54 torch
55 ];
56 pytestFlagsArray = [ "tests" ];
57 # don't require PaddlePaddle (not in Nixpkgs), Flax, or Tensorflow (onerous) to run tests:
58 disabledTestPaths =
59 [
60 "tests/test_flax_comparison.py"
61 "tests/test_paddle_comparison.py"
62 "tests/test_tf_comparison.py"
63 ]
64 ++ lib.optionals stdenv.isDarwin [
65 # don't require mlx (not in Nixpkgs) to run tests
66 "tests/test_mlx_comparison.py"
67 ];
68
69 pythonImportsCheck = [ "safetensors" ];
70
71 meta = with lib; {
72 homepage = "https://github.com/huggingface/safetensors";
73 description = "Fast (zero-copy) and safe (unlike pickle) format for storing tensors";
74 changelog = "https://github.com/huggingface/safetensors/releases/tag/v${version}";
75 license = licenses.asl20;
76 maintainers = with maintainers; [ bcdarwin ];
77 };
78}