pkgs/development/python-modules/sparsezoo/default.nix

   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
  27 buildPythonPackage 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 }