fix ultravox model installations

stikkireddy · Sep 20, 2024 · 6696c5f · 6696c5f
1 parent 02083ad
commit 6696c5f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/mlflow_extensions/databricks/prebuilt/audio/vllm.py b/mlflow_extensions/databricks/prebuilt/audio/vllm.py
@@ -15,6 +15,9 @@
         trust_remote_code=True,
         guided_decoding_backend="outlines",
         max_model_len=64000,
+        library_overrides={
+            "vllm": '"vllm[audio]==0.6.0"'
+        }
     ),
     serving_config=ServingConfig(
         # rough estimate for the engines this includes model weights + kv cache + overhead + intermediate states
@@ -30,6 +33,9 @@
         model="fixie-ai/ultravox-v0_4",
         trust_remote_code=True,
         guided_decoding_backend="outlines",
+        library_overrides={
+            "vllm": 'vllm[audio]==0.6.0'
+        }
     ),
     serving_config=ServingConfig(
         # rough estimate for the engines this includes model weights + kv cache + overhead + intermediate states