tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
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.25.3";
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-SVyGtWFI5GjfxbaVKICf+QSSMYI62dDVMzphu8TngvY=";
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 ] ++ transformers.optional-dependencies.sentencepiece;
54
55 optional-dependencies = {
56 onnxruntime = [
57 onnx
58 datasets
59 protobuf
60 onnxruntime
61 ];
62 exporters = [
63 onnx
64 timm
65 onnxruntime
66 protobuf
67 ];
68 exporters-tf = [
69 onnx
70 timm
71 h5py
72 tf2onnx
73 onnxruntime
74 numpy
75 datasets
76 tensorflow
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}