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