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