python310Packages.onnxruntime-tools: init at 1.7.0

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

@@ -0,0 +1,50 @@
+{ lib
+, buildPythonPackage
+, fetchPypi
+, coloredlogs
+, numpy
+, onnx
+, packaging
+, psutil
+, py-cpuinfo
+, py3nvml
+, sympy
+}:
+
+buildPythonPackage rec {
+  pname = "onnxruntime-tools";
+  version = "1.7.0";
+  format = "wheel";
+
+  # the build distribution doesn't work at all, it seems to expect the same structure
+  # as the github source repo.
+  # The github source wasn't immediately obvious how to build for this subpackage.
+  src = fetchPypi {
+    pname = "onnxruntime_tools";
+    inherit version;
+    format = "wheel";
+    dist = "py3";
+    python = "py3";
+    hash = "sha256-Hf+Ii1xIKsW8Yn8S4QhEX+/LPWAMQ/Y2M5dTFv5hetg=";
+  };
+
+  propagatedBuildInputs = [
+    coloredlogs
+    numpy
+    onnx
+    packaging
+    psutil
+    py-cpuinfo
+    py3nvml
+    sympy
+  ];
+
+  pythonImportsCheck = [ "onnxruntime_tools" ];
+
+  meta = with lib; {
+    description = "Transformers Model Optimization Tool of ONNXRuntime";
+    homepage = "https://pypi.org/project/onnxruntime-tools/";
+    license = with licenses; [ mit ];
+    maintainers = with maintainers; [ happysalada ];
+  };
+}

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

@@ -17,6 +17,7 @@
 # optional dependencies
 , scikit-learn
 , tensorflow
+, onnxconverter-common
 , torch
 , accelerate
 , faiss
@@ -27,6 +28,7 @@
 , optax
 , ftfy
 , onnxruntime
+, onnxruntime-tools
 , cookiecutter
 , sagemaker
 , fairscale
@@ -96,7 +98,7 @@ buildPythonPackage rec {
     ];
     tf = [
       tensorflow
-      # onnxconverter-common
+      onnxconverter-common
       # tf2onnx
       # tensorflow-text
       # keras-nlp
@@ -113,10 +115,10 @@ buildPythonPackage rec {
     ftfy = [ ftfy ];
     onnxruntime = [
       onnxruntime
-      # onnxruntime-tools
+      onnxruntime-tools
     ];
     onnx = [
-      # onnxconverter-common
+      onnxconverter-common
       # tf2onnx
       onnxruntime
     ];

pkgs/top-level/python-packages.nix

@@ -7228,6 +7228,8 @@ self: super: with self; {
     };
   };
 
+  onnxruntime-tools = callPackage ../development/python-modules/onnxruntime-tools { };
+
   onvif-zeep-async = callPackage ../development/python-modules/onvif-zeep-async { };
 
   oocsi = callPackage ../development/python-modules/oocsi { };