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