pkgs/development/python-modules/openllm/default.nix

   1 { lib
   2 , buildPythonPackage
   3 , hatch-fancy-pypi-readme
   4 , hatch-vcs
   5 , hatchling
   6 , pytestCheckHook
   7 , pythonOlder
   8 , pythonRelaxDepsHook
   9 , accelerate
  10 , bentoml
  11 , bitsandbytes
  12 , click
  13 , datasets
  14 , docker
  15 , einops
  16 , fairscale
  17 , flax
  18 , ghapi
  19 , hypothesis
  20 , ipython
  21 , jax
  22 , jaxlib
  23 , jupyter
  24 , jupytext
  25 , keras
  26 , nbformat
  27 , notebook
  28 , openai
  29 , openllm-client
  30 , openllm-core
  31 , optimum
  32 , peft
  33 , pytest-mock
  34 , pytest-randomly
  35 , pytest-rerunfailures
  36 , pytest-xdist
  37 , ray
  38 , safetensors
  39 , scipy
  40 , sentencepiece
  41 , soundfile
  42 , syrupy
  43 , tabulate
  44 , tensorflow
  45 , tiktoken
  46 , transformers
  47 , openai-triton
  48 , xformers
  49 }:
  50
  51 buildPythonPackage rec {
  52   inherit (openllm-core) src version;
  53   pname = "openllm";
  54   pyproject = true;
  55
  56   disabled = pythonOlder "3.8";
  57
  58   sourceRoot = "source/openllm-python";
  59
  60   nativeBuildInputs = [
  61     hatch-fancy-pypi-readme
  62     hatch-vcs
  63     hatchling
  64     pythonRelaxDepsHook
  65   ];
  66
  67   pythonRemoveDeps = [
  68     # remove cuda-python as it has an unfree license
  69     "cuda-python"
  70   ];
  71
  72   propagatedBuildInputs = [
  73     accelerate
  74     bentoml
  75     bitsandbytes
  76     click
  77     ghapi
  78     openllm-client
  79     openllm-core
  80     optimum
  81     safetensors
  82     tabulate
  83     transformers
  84   ] ++ bentoml.optional-dependencies.io
  85   ++ tabulate.optional-dependencies.widechars
  86   ++ transformers.optional-dependencies.tokenizers
  87   ++ transformers.optional-dependencies.torch;
  88
  89   passthru.optional-dependencies = {
  90     agents = [
  91       # diffusers
  92       soundfile
  93       transformers
  94     ] ++ transformers.optional-dependencies.agents;
  95     baichuan = [
  96       # cpm-kernels
  97       sentencepiece
  98     ];
  99     chatglm = [
 100       # cpm-kernels
 101       sentencepiece
 102     ];
 103     falcon = [
 104       einops
 105       xformers
 106     ];
 107     fine-tune = [
 108       accelerate
 109       bitsandbytes
 110       datasets
 111       peft
 112       # trl
 113     ];
 114     flan-t5 = [
 115       flax
 116       jax
 117       jaxlib
 118       keras
 119       tensorflow
 120     ];
 121     ggml = [
 122       # ctransformers
 123     ];
 124     gptq = [
 125       # auto-gptq
 126       optimum
 127     ]; # ++ autogptq.optional-dependencies.triton;
 128     grpc = [
 129       openllm-client
 130     ] ++ openllm-client.optional-dependencies.grpc;
 131     llama = [
 132       fairscale
 133       sentencepiece
 134       scipy
 135     ];
 136     mpt = [
 137       einops
 138       openai-triton
 139     ];
 140     openai = [
 141       openai
 142       tiktoken
 143     ] ++ openai.optional-dependencies.embeddings;
 144     opt = [
 145       flax
 146       jax
 147       jaxlib
 148       keras
 149       tensorflow
 150     ];
 151     playground = [
 152       ipython
 153       jupyter
 154       jupytext
 155       nbformat
 156       notebook
 157     ];
 158     starcoder = [
 159       bitsandbytes
 160     ];
 161     vllm = [
 162       ray
 163       # vllm
 164     ];
 165     full = with passthru.optional-dependencies; (
 166       agents ++ baichuan ++ chatglm ++ falcon ++ fine-tune ++ flan-t5 ++ ggml ++ gptq ++ llama ++ mpt ++ openai ++ opt ++ playground ++ starcoder ++ vllm
 167     );
 168     all = passthru.optional-dependencies.full;
 169   };
 170
 171   nativeCheckInputs = [
 172     docker
 173     hypothesis
 174     pytest-mock
 175     pytest-randomly
 176     pytest-rerunfailures
 177     pytest-xdist
 178     pytestCheckHook
 179     syrupy
 180   ];
 181
 182   preCheck = ''
 183     export HOME=$TMPDIR
 184     # skip GPUs test on CI
 185     export GITHUB_ACTIONS=1
 186     # disable hypothesis' deadline
 187     export CI=1
 188   '';
 189
 190   disabledTests = [
 191     # these tests access to huggingface.co
 192     "test_opt_125m"
 193     "test_opt_125m"
 194     "test_flan_t5"
 195     "test_flan_t5"
 196   ];
 197
 198   pythonImportsCheck = [ "openllm" ];
 199
 200   meta = with lib; {
 201     description = "Operating LLMs in production";
 202     homepage = "https://github.com/bentoml/OpenLLM/tree/main/openllm-python";
 203     changelog = "https://github.com/bentoml/OpenLLM/blob/${src.rev}/CHANGELOG.md";
 204     license = licenses.asl20;
 205     maintainers = with maintainers; [ happysalada natsukium ];
 206   };
 207 }