pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, pythonOlder
5, coloredlogs
6, datasets
7, evaluate
8, h5py
9, huggingface-hub
10, numpy
11, onnx
12, onnxruntime
13, packaging
14, protobuf
15, sympy
16, tensorflow
17, tf2onnx
18, timm
19, torch
20, transformers
21}:
22
23buildPythonPackage rec {
24 pname = "optimum";
25 version = "1.14.0";
26 format = "setuptools";
27
28 disabled = pythonOlder "3.7";
29
30 src = fetchFromGitHub {
31 owner = "huggingface";
32 repo = "optimum";
33 rev = "refs/tags/v${version}";
34 hash = "sha256-fCLMDbWfEAG7pTFKWv67fpr6lNTxRTCBkppgDZR51eY=";
35 };
36
37 propagatedBuildInputs = [
38 coloredlogs
39 datasets
40 huggingface-hub
41 numpy
42 packaging
43 sympy
44 torch
45 transformers
46 ] ++ transformers.optional-dependencies.sentencepiece;
47
48 passthru.optional-dependencies = {
49 onnxruntime = [
50 onnx
51 onnxruntime
52 datasets
53 evaluate
54 protobuf
55 ];
56 exporters = [
57 onnx
58 onnxruntime
59 timm
60 ];
61 exporters-tf = [
62 tensorflow
63 tf2onnx
64 onnx
65 onnxruntime
66 timm
67 h5py
68 numpy
69 ];
70 diffusers = [
71 # diffusers
72 ];
73 intel = [
74 # optimum-intel
75 ];
76 openvino = [
77 # optimum-intel
78 ]; # ++ optimum-intel.optional-dependencies.openvino;
79 nncf = [
80 # optimum-intel
81 ]; # ++ optimum-intel.optional-dependencies.nncf;
82 neural-compressor = [
83 # optimum-intel
84 ]; # ++ optimum-intel.optional-dependencies.neural-compressor;
85 graphcore = [
86 # optimum-graphcore
87 ];
88 habana = [
89 transformers
90 # optimum-habana
91 ];
92 neuron = [
93 # optimum-neuron
94 ]; # ++ optimum-neuron.optional-dependencies.neuron;
95 neuronx = [
96 # optimum-neuron
97 ]; # ++ optimum-neuron.optional-dependencies.neuronx;
98 furiosa = [
99 # optimum-furiosa
100 ];
101 };
102
103 # almost all tests try to connect to https://huggingface.co
104 doCheck = false;
105
106 pythonImportsCheck = [ "optimum" ];
107
108 meta = with lib; {
109 description = "Accelerate training and inference of 🤗 Transformers and 🤗 Diffusers with easy to use hardware optimization tools";
110 homepage = "https://github.com/huggingface/optimum";
111 changelog = "https://github.com/huggingface/optimum/releases/tag/${src.rev}";
112 license = licenses.asl20;
113 maintainers = with maintainers; [ natsukium ];
114 };
115}