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

   1 { lib
   2 , buildPythonPackage
   3 , autoPatchelfHook
   4 , pythonRelaxDepsHook
   5 , onnxruntime
   6 , coloredlogs
   7 , numpy
   8 , packaging
   9 , oneDNN
  10
  11 }:
  12
  13 # onnxruntime requires an older protobuf.
  14 # Doing an override in protobuf in the python-packages set
  15 # can give you a functioning Python package but note not
  16 # all Python packages will be compatible then.
  17 #
  18 # Because protobuf is not always needed we remove it
  19 # as a runtime dependency from our wheel.
  20 #
  21 # We do include here the non-Python protobuf so the shared libs
  22 # link correctly. If you do also want to include the Python
  23 # protobuf, you can add it to your Python env, but be aware
  24 # the version likely mismatches with what is used here.
  25
  26 buildPythonPackage {
  27   inherit (onnxruntime) pname version;
  28   format = "wheel";
  29   src = onnxruntime.dist;
  30
  31   unpackPhase = ''
  32     cp -r $src dist
  33     chmod +w dist
  34   '';
  35
  36   nativeBuildInputs = [
  37     autoPatchelfHook
  38     pythonRelaxDepsHook
  39   ];
  40
  41   # This project requires fairly large dependencies such as sympy which we really don't always need.
  42   pythonRemoveDeps = [
  43     "flatbuffers"
  44     "protobuf"
  45     "sympy"
  46   ];
  47
  48   # Libraries are not linked correctly.
  49   buildInputs = [
  50     oneDNN
  51     onnxruntime.protobuf
  52   ];
  53
  54   propagatedBuildInputs = [
  55     coloredlogs
  56     # flatbuffers
  57     numpy
  58     packaging
  59     # protobuf
  60     # sympy
  61   ];
  62
  63   meta = onnxruntime.meta // { maintainers = with lib.maintainers; [ fridh ]; };
  64 }