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

   1 { lib
   2 , buildPythonPackage
   3 , fetchFromGitHub
   4 , isPy27
   5 , numpy
   6 , scikitlearn
   7 , pytestCheckHook
   8 , pytorch
   9 , torchvision
  10 , tqdm
  11 }:
  12
  13 buildPythonPackage rec {
  14   pname   = "pytorch-metric-learning";
  15   version = "0.9.95";
  16
  17   disabled = isPy27;
  18
  19   src = fetchFromGitHub {
  20     owner = "KevinMusgrave";
  21     repo = pname;
  22     rev = "v${version}";
  23     sha256 = "1msvs1j3n47762ahm21bnkk2qqabxw8diiyi7s420x4zg24mr23g";
  24   };
  25
  26   propagatedBuildInputs = [
  27     numpy
  28     pytorch
  29     scikitlearn
  30     torchvision
  31     tqdm
  32   ];
  33
  34   preCheck = ''
  35     export HOME=$TMP
  36     export TEST_DEVICE=cpu
  37     export TEST_DTYPES=float32,float64  # half-precision tests fail on CPU
  38   '';
  39   # package only requires `unittest`, but use `pytest` to exclude tests
  40   checkInputs = [ pytestCheckHook ];
  41   disabledTests = [
  42     # requires FAISS (not in Nixpkgs)
  43     "test_accuracy_calculator_and_faiss"
  44     # require network access:
  45     "test_get_nearest_neighbors"
  46     "test_tuplestoweights_sampler"
  47     "test_untrained_indexer"
  48   ];
  49
  50   meta = {
  51     description = "Metric learning library for PyTorch";
  52     homepage = "https://github.com/KevinMusgrave/pytorch-metric-learning";
  53     changelog = "https://github.com/KevinMusgrave/pytorch-metric-learning/releases/tag/v${version}";
  54     license = lib.licenses.mit;
  55     maintainers = with lib.maintainers; [ bcdarwin ];
  56   };
  57 }