1 files changed, 26 insertions, 43 deletions
diff --git a/pkgs/development/python-modules/openllm/default.nix b/pkgs/development/python-modules/openllm/default.nix
index b9f3d2b6fa3b5..b1cad320a11d3 100644
--- a/pkgs/development/python-modules/openllm/default.nix
+++ b/pkgs/development/python-modules/openllm/default.nix
@@ -9,20 +9,18 @@
 , accelerate
 , bentoml
 , bitsandbytes
+, build
 , click
+, ctranslate2
 , datasets
 , docker
 , einops
-, fairscale
-, flax
 , ghapi
+, huggingface-hub
 , hypothesis
 , ipython
-, jax
-, jaxlib
 , jupyter
 , jupytext
-, keras
 , nbformat
 , notebook
 , openai
@@ -34,14 +32,12 @@
 , pytest-randomly
 , pytest-rerunfailures
 , pytest-xdist
-, ray
 , safetensors
 , scipy
 , sentencepiece
 , soundfile
 , syrupy
 , tabulate
-, tensorflow
 , tiktoken
 , transformers
 , openai-triton
@@ -73,12 +69,16 @@ buildPythonPackage rec {
     accelerate
     bentoml
     bitsandbytes
+    build
     click
+    einops
     ghapi
     openllm-client
     openllm-core
     optimum
     safetensors
+    scipy
+    sentencepiece
     tabulate
     transformers
   ] ++ bentoml.optional-dependencies.io
@@ -92,62 +92,43 @@ buildPythonPackage rec {
       soundfile
       transformers
     ] ++ transformers.optional-dependencies.agents;
+    awq = [
+      # autoawq
+    ];
     baichuan = [
       # cpm-kernels
-      sentencepiece
     ];
     chatglm = [
       # cpm-kernels
-      sentencepiece
+    ];
+    ctranslate = [
+      ctranslate2
     ];
     falcon = [
-      einops
       xformers
     ];
     fine-tune = [
-      accelerate
-      bitsandbytes
       datasets
+      huggingface-hub
       peft
       # trl
     ];
-    flan-t5 = [
-      flax
-      jax
-      jaxlib
-      keras
-      tensorflow
-    ];
     ggml = [
       # ctransformers
     ];
     gptq = [
       # auto-gptq
-      optimum
     ]; # ++ autogptq.optional-dependencies.triton;
     grpc = [
-      openllm-client
-    ] ++ openllm-client.optional-dependencies.grpc;
-    llama = [
-      fairscale
-      sentencepiece
-      scipy
-    ];
+      bentoml
+    ] ++ bentoml.optional-dependencies.grpc;
     mpt = [
-      einops
       openai-triton
     ];
     openai = [
       openai
       tiktoken
-    ] ++ openai.optional-dependencies.embeddings;
-    opt = [
-      flax
-      jax
-      jaxlib
-      keras
-      tensorflow
-    ];
+    ] ++ openai.optional-dependencies.datalib;
     playground = [
       ipython
       jupyter
@@ -159,11 +140,10 @@ buildPythonPackage rec {
       bitsandbytes
     ];
     vllm = [
-      ray
       # vllm
     ];
     full = with passthru.optional-dependencies; (
-      agents ++ baichuan ++ chatglm ++ falcon ++ fine-tune ++ flan-t5 ++ ggml ++ gptq ++ llama ++ mpt ++ openai ++ opt ++ playground ++ starcoder ++ vllm
+      agents ++ awq ++ baichuan ++ chatglm ++ ctranslate ++ falcon ++ fine-tune ++ ggml ++ gptq ++ mpt ++ openai ++ playground ++ starcoder ++ vllm
     );
     all = passthru.optional-dependencies.full;
   };
@@ -187,12 +167,15 @@ buildPythonPackage rec {
     export CI=1
   '';
 
+  disabledTestPaths = [
+    # require network access
+    "tests/models"
+  ];
+
   disabledTests = [
-    # these tests access to huggingface.co
-    "test_opt_125m"
-    "test_opt_125m"
-    "test_flan_t5"
-    "test_flan_t5"
+    # incompatible with recent TypedDict
+    # https://github.com/bentoml/OpenLLM/blob/f3fd32d596253ae34c68e2e9655f19f40e05f666/openllm-python/tests/configuration_test.py#L18-L21
+    "test_missing_default"
   ];
 
   pythonImportsCheck = [ "openllm" ];