tjh.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 setuptools,
7 pydantic,
8 torch,
9 transformers,
10 nbconvert,
11 nbformat,
12 pytestCheckHook,
13}:
14
15buildPythonPackage rec {
16 pname = "compressed-tensors";
17 version = "0.9.2";
18 pyproject = true;
19
20 # Release on PyPI is missing the `utils` directory, which `setup.py` wants to import
21 src = fetchFromGitHub {
22 owner = "neuralmagic";
23 repo = "compressed-tensors";
24 tag = version;
25 hash = "sha256-PxW8zseDUF0EOh7E/N8swwgFTfvkoTpp+d3ngAUpFNU=";
26 };
27
28 build-system = [ setuptools ];
29
30 dependencies = [
31 pydantic
32 torch
33 transformers
34 ];
35
36 doCheck = true;
37
38 pythonImportsCheck = [ "compressed_tensors" ];
39
40 nativeCheckInputs = [
41 nbconvert
42 nbformat
43 pytestCheckHook
44 ];
45
46 disabledTests = [
47 # these try to download models from HF Hub
48 "test_get_observer_token_count"
49 "test_kv_cache_quantization"
50 "test_target_prioritization"
51 "test_load_compressed_sharded"
52 "test_save_compressed_model"
53 "test_apply_tinyllama_dynamic_activations"
54 ];
55
56 disabledTestPaths = [
57 # these try to download models from HF Hub
58 "tests/test_quantization/lifecycle/test_apply.py"
59 ];
60
61 meta = {
62 description = "Safetensors extension to efficiently store sparse quantized tensors on disk";
63 homepage = "https://github.com/neuralmagic/compressed-tensors";
64 changelog = "https://github.com/neuralmagic/compressed-tensors/releases/tag/${version}";
65 license = lib.licenses.asl20;
66 maintainers = [ ];
67 };
68}