pkgs/development/python-modules/pytorch-metric-learning/default.nix

   1 {
   2   stdenv,
   3   lib,
   4   buildPythonPackage,
   5   fetchFromGitHub,
   6   isPy27,
   7   config,
   8
   9   # build-system
  10   setuptools,
  11
  12   # dependencies
  13   numpy,
  14   scikit-learn,
  15   torch,
  16   tqdm,
  17
  18   # optional-dependencies
  19   faiss,
  20   tensorboard,
  21
  22   # tests
  23   cudaSupport ? config.cudaSupport,
  24   pytestCheckHook,
  25   torchvision
  26 }:
  27
  28 buildPythonPackage rec {
  29   pname = "pytorch-metric-learning";
  30   version = "2.7.0";
  31   pyproject = true;
  32
  33   disabled = isPy27;
  34
  35   src = fetchFromGitHub {
  36     owner = "KevinMusgrave";
  37     repo = pname;
  38     rev = "refs/tags/v${version}";
  39     hash = "sha256-mxAl4GMyAWtvocc68Ac3z1+W13k9OOK7aQFfB7X0f9c=";
  40   };
  41
  42   build-system = [
  43     setuptools
  44   ];
  45
  46   dependencies = [
  47     numpy
  48     torch
  49     scikit-learn
  50     tqdm
  51   ];
  52
  53   optional-dependencies = {
  54     with-hooks = [
  55       # TODO: record-keeper
  56       faiss
  57       tensorboard
  58     ];
  59     with-hooks-cpu = [
  60       # TODO: record-keeper
  61       faiss
  62       tensorboard
  63     ];
  64   };
  65
  66   preCheck = ''
  67     export HOME=$TMP
  68     export TEST_DEVICE=cpu
  69     export TEST_DTYPES=float32,float64  # half-precision tests fail on CPU
  70   '';
  71
  72   # package only requires `unittest`, but use `pytest` to exclude tests
  73   nativeCheckInputs = [
  74     pytestCheckHook
  75     torchvision
  76   ] ++ lib.flatten (lib.attrValues optional-dependencies);
  77
  78   disabledTests = [
  79     # network access
  80     "test_tuplestoweights_sampler"
  81     "test_metric_loss_only"
  82     "test_add_to_indexer"
  83     "test_get_nearest_neighbors"
  84     "test_list_of_text"
  85     "test_untrained_indexer"
  86   ] ++ lib.optionals stdenv.hostPlatform.isDarwin [
  87     # AttributeError: module 'torch.distributed' has no attribute 'init_process_group'
  88     "test_single_proc"
  89   ] ++ lib.optionals cudaSupport [
  90     # crashes with SIGBART
  91     "test_accuracy_calculator_and_faiss_with_torch_and_numpy"
  92     "test_accuracy_calculator_large_k"
  93     "test_custom_knn"
  94     "test_global_embedding_space_tester"
  95     "test_global_two_stream_embedding_space_tester"
  96     "test_index_type"
  97     "test_k_warning"
  98     "test_many_tied_distances"
  99     "test_query_within_reference"
 100     "test_tied_distances"
 101     "test_with_same_parent_label_tester"
 102   ];
 103
 104   meta = {
 105     description = "Metric learning library for PyTorch";
 106     homepage = "https://github.com/KevinMusgrave/pytorch-metric-learning";
 107     changelog = "https://github.com/KevinMusgrave/pytorch-metric-learning/releases/tag/v${version}";
 108     license = lib.licenses.mit;
 109     maintainers = with lib.maintainers; [ bcdarwin ];
 110   };
 111 }