tjh.dev
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 pyyaml,
15 requests,
16 tqdm,
17 pandas,
18 protobuf,
19 py-machineid,
20 pydantic,
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 rev = "refs/tags/v${version}";
36 hash = "sha256-c4F95eVvj673eFO/rbmv4LY3pGmqo+arbsYqElznwdA=";
37 };
38
39 build-system = [ setuptools ];
40
41 pythonRelaxDeps = [ "onnx" ];
42
43 dependencies = [
44 click
45 geocoder
46 numpy
47 onnx
48 pyyaml
49 requests
50 tqdm
51 pandas
52 protobuf
53 py-machineid
54 pydantic
55 ];
56
57 pythonImportsCheck = [ "sparsezoo" ];
58
59 nativeCheckInputs = [
60 matplotlib
61 pytestCheckHook
62 ];
63
64 disabledTests = [
65 # Require network access
66 "test_analysis"
67 "test_custom_default"
68 "test_draw_operation_chart"
69 "test_draw_parameter_chart"
70 "test_draw_parameter_operation_combined_chart"
71 "test_draw_sparsity_by_layer_chart"
72 "test_extract_node_id"
73 "test_fail_default_on_empty"
74 "test_folder_structure"
75 "test_from_model_analysis"
76 "test_generate_outputs"
77 "test_get_layer_and_op_counts"
78 "test_get_node_bias"
79 "test_get_node_bias_name"
80 "test_get_node_four_block_sparsity"
81 "test_get_node_num_four_block_zeros_and_size"
82 "test_get_node_num_zeros_and_size"
83 "test_get_node_sparsity"
84 "test_get_node_weight_and_shape"
85 "test_get_node_weight_name"
86 "test_get_ops_dict"
87 "test_get_zero_point"
88 "test_graphql_api_response"
89 "test_is_four_block_sparse_layer"
90 "test_is_parameterized_prunable_layer"
91 "test_is_quantized_layer"
92 "test_is_sparse_layer"
93 "test_latency_extractor"
94 "test_load_files_from_stub"
95 "test_model_from_stub"
96 "test_model_gz_extraction_from_local_files"
97 "test_model_gz_extraction_from_stub"
98 "test_recipe_getters"
99 "test_search_models"
100 "test_setup_model_from_objects"
101 "test_setup_model_from_paths"
102 "test_validate"
103 ];
104
105 disabledTestPaths = [
106 # Require network access
107 "tests/sparsezoo/analyze/test_model_analysis_creation.py"
108 "tests/sparsezoo/deployment_package/utils/test_utils.py"
109 ];
110
111 meta = {
112 description = "Neural network model repository for highly sparse and sparse-quantized models with matching sparsification recipes";
113 homepage = "https://github.com/neuralmagic/sparsezoo";
114 changelog = "https://github.com/neuralmagic/sparsezoo/releases/tag/v${version}";
115 license = lib.licenses.asl20;
116 maintainers = with lib.maintainers; [ GaetanLepage ];
117 };
118}