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