pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 setuptools,
8 setuptools-scm,
9
10 # dependencies
11 frozendict,
12 pydantic,
13 torch,
14 transformers,
15
16 # tests
17 nbconvert,
18 nbformat,
19 pytestCheckHook,
20}:
21
22buildPythonPackage rec {
23 pname = "compressed-tensors";
24 version = "0.11.0";
25 pyproject = true;
26
27 # Release on PyPI is missing the `utils` directory, which `setup.py` wants to import
28 src = fetchFromGitHub {
29 owner = "neuralmagic";
30 repo = "compressed-tensors";
31 tag = version;
32 hash = "sha256-sSXn4/N/Pn+wOCY1Z0ziqFxfMRvRA1c90jPOBe+SwZw=";
33 };
34
35 postPatch = ''
36 substituteInPlace pyproject.toml \
37 --replace-fail "setuptools_scm==8.2.0" "setuptools_scm"
38 '';
39
40 build-system = [
41 setuptools
42 setuptools-scm
43 ];
44
45 dependencies = [
46 frozendict
47 pydantic
48 torch
49 transformers
50 ];
51
52 doCheck = true;
53
54 pythonImportsCheck = [ "compressed_tensors" ];
55
56 nativeCheckInputs = [
57 nbconvert
58 nbformat
59 pytestCheckHook
60 ];
61
62 disabledTests = [
63 # these try to download models from HF Hub
64 "test_apply_tinyllama_dynamic_activations"
65 "test_compress_model"
66 "test_compress_model_meta"
67 "test_compressed_linear_from_linear_usage"
68 "test_decompress_model"
69 "test_get_observer_token_count"
70 "test_kv_cache_quantization"
71 "test_load_compressed_sharded"
72 "test_model_forward_pass"
73 "test_save_compressed_model"
74 "test_target_prioritization"
75 ];
76
77 disabledTestPaths = [
78 # these try to download models from HF Hub
79 "tests/test_quantization/lifecycle/test_apply.py"
80 ];
81
82 meta = {
83 description = "Safetensors extension to efficiently store sparse quantized tensors on disk";
84 homepage = "https://github.com/neuralmagic/compressed-tensors";
85 changelog = "https://github.com/neuralmagic/compressed-tensors/releases/tag/${src.tag}";
86 license = lib.licenses.asl20;
87 maintainers = [ ];
88 };
89}