1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 setuptools,
8
9 # dependencies
10 click,
11 geocoder,
12 numpy,
13 onnx,
14 pandas,
15 protobuf,
16 py-machineid,
17 pydantic,
18 pyyaml,
19 requests,
20 tqdm,
21
22 # checks
23 matplotlib,
24 pytestCheckHook,
25}:
26
27buildPythonPackage rec {
28 pname = "sparsezoo";
29 version = "1.8.1";
30 pyproject = true;
31
32 src = fetchFromGitHub {
33 owner = "neuralmagic";
34 repo = "sparsezoo";
35 tag = "v${version}";
36 hash = "sha256-c4F95eVvj673eFO/rbmv4LY3pGmqo+arbsYqElznwdA=";
37 };
38
39 build-system = [ setuptools ];
40
41 pythonRelaxDeps = [
42 "numpy"
43 "onnx"
44 ];
45
46 dependencies = [
47 click
48 geocoder
49 numpy
50 onnx
51 pandas
52 protobuf
53 py-machineid
54 pydantic
55 pyyaml
56 requests
57 tqdm
58 ];
59
60 pythonImportsCheck = [ "sparsezoo" ];
61
62 nativeCheckInputs = [
63 matplotlib
64 pytestCheckHook
65 ];
66
67 disabledTests = [
68 # Require network access
69 "test_analysis"
70 "test_custom_default"
71 "test_draw_operation_chart"
72 "test_draw_parameter_chart"
73 "test_draw_parameter_operation_combined_chart"
74 "test_draw_sparsity_by_layer_chart"
75 "test_extract_node_id"
76 "test_fail_default_on_empty"
77 "test_folder_structure"
78 "test_from_model_analysis"
79 "test_generate_outputs"
80 "test_get_layer_and_op_counts"
81 "test_get_node_bias"
82 "test_get_node_bias_name"
83 "test_get_node_four_block_sparsity"
84 "test_get_node_num_four_block_zeros_and_size"
85 "test_get_node_num_zeros_and_size"
86 "test_get_node_sparsity"
87 "test_get_node_weight_and_shape"
88 "test_get_node_weight_name"
89 "test_get_ops_dict"
90 "test_get_zero_point"
91 "test_graphql_api_response"
92 "test_is_four_block_sparse_layer"
93 "test_is_parameterized_prunable_layer"
94 "test_is_quantized_layer"
95 "test_is_sparse_layer"
96 "test_latency_extractor"
97 "test_load_files_from_stub"
98 "test_model_from_stub"
99 "test_model_gz_extraction_from_local_files"
100 "test_model_gz_extraction_from_stub"
101 "test_recipe_getters"
102 "test_search_models"
103 "test_setup_model_from_objects"
104 "test_setup_model_from_paths"
105 "test_validate"
106 ];
107
108 disabledTestPaths = [
109 # Require network access
110 "tests/sparsezoo/analyze/test_model_analysis_creation.py"
111 "tests/sparsezoo/deployment_package/utils/test_utils.py"
112 ];
113
114 meta = {
115 description = "Neural network model repository for highly sparse and sparse-quantized models with matching sparsification recipes";
116 homepage = "https://github.com/neuralmagic/sparsezoo";
117 changelog = "https://github.com/neuralmagic/sparsezoo/releases/tag/v${version}";
118 license = lib.licenses.asl20;
119 maintainers = with lib.maintainers; [ GaetanLepage ];
120 };
121}