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