tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 pythonAtLeast,
7
8 # build-system
9 setuptools,
10 setuptools-scm,
11 llvmPackages,
12
13 # dependencies
14 frozendict,
15 loguru,
16 pydantic,
17 torch,
18 transformers,
19
20 # tests
21 nbconvert,
22 nbformat,
23 pytestCheckHook,
24}:
25
26buildPythonPackage (finalAttrs: {
27 pname = "compressed-tensors";
28 version = "0.14.0.1";
29 pyproject = true;
30
31 # Release on PyPI is missing the `utils` directory, which `setup.py` wants to import
32 src = fetchFromGitHub {
33 owner = "neuralmagic";
34 repo = "compressed-tensors";
35 tag = finalAttrs.version;
36 hash = "sha256-bncglC6qRE4D6rjoNOXRKh+7ZMSTc9zc3NVzA+NaXuE=";
37 };
38
39 postPatch = ''
40 substituteInPlace pyproject.toml \
41 --replace-fail "setuptools_scm==8.2.0" "setuptools_scm"
42 '';
43
44 build-system = [
45 setuptools
46 setuptools-scm
47 ];
48
49 buildInputs = lib.optional stdenv.cc.isClang llvmPackages.openmp;
50
51 dependencies = [
52 frozendict
53 loguru
54 pydantic
55 torch
56 transformers
57 ];
58
59 pythonRelaxDeps = [ "transformers" ];
60
61 doCheck = true;
62
63 pythonImportsCheck = [ "compressed_tensors" ];
64
65 nativeCheckInputs = [
66 nbconvert
67 nbformat
68 pytestCheckHook
69 ];
70
71 disabledTests = [
72 # these try to download models from HF Hub
73 "test_apply_tinyllama_dynamic_activations"
74 "test_compress_model"
75 "test_compress_model_meta"
76 "test_compressed_linear_from_linear_usage"
77 "test_decompress_model"
78 "test_get_observer_token_count"
79 "test_kv_cache_quantization"
80 "test_load_compressed_sharded"
81 "test_model_forward_pass"
82 "test_save_compressed_model"
83 "test_target_prioritization"
84 "test_expand_targets_with_llama_stories"
85
86 # AssertionError: Torch not compiled with CUDA enabled
87 "test_register_parameter"
88 "test_register_parameter_invalidates"
89 "test_set_item"
90 "test_set_item_buffers"
91 "test_update_offload_parameter_with_grad"
92 ];
93
94 disabledTestPaths = [
95 # these try to download models from HF Hub
96 "tests/test_quantization/lifecycle/test_apply.py"
97 # RuntimeError: The weights trying to be saved contained shared tensors
98 "tests/test_transform/factory/test_serialization.py::test_serialization[True-hadamard]"
99 "tests/test_transform/factory/test_serialization.py::test_serialization[True-random-hadamard]"
100 "tests/test_transform/factory/test_serialization.py::test_serialization[False-hadamard]"
101 "tests/test_transform/factory/test_serialization.py::test_serialization[False-random-hadamard]"
102 "tests/test_transform/factory/test_serialization.py::test_serialization[False-True-hadamard]"
103 "tests/test_transform/factory/test_serialization.py::test_serialization[False-True-random-hadamard]"
104 "tests/test_transform/factory/test_serialization.py::test_serialization[False-False-hadamard]"
105 "tests/test_transform/factory/test_serialization.py::test_serialization[False-False-random-hadamard]"
106
107 # AssertionError: Torch not compiled with CUDA enabled
108 "tests/test_transform/factory/test_serialization.py::test_serialization[True-True-hadamard]"
109 "tests/test_transform/factory/test_serialization.py::test_serialization[True-True-random-hadamard]"
110 "tests/test_transform/factory/test_serialization.py::test_serialization[True-False-hadamard]"
111 "tests/test_transform/factory/test_serialization.py::test_serialization[True-False-random-hadamard]"
112 ];
113
114 meta = {
115 description = "Safetensors extension to efficiently store sparse quantized tensors on disk";
116 homepage = "https://github.com/neuralmagic/compressed-tensors";
117 changelog = "https://github.com/neuralmagic/compressed-tensors/releases/tag/${finalAttrs.src.tag}";
118 license = lib.licenses.asl20;
119 maintainers = with lib.maintainers; [ sarahec ];
120 };
121})