OAI: Add draft model dir to inline load

Was not pushed before and caused errors of the kwargs being None. Signed-off-by: kingbri <[email protected]>
theroyallab · Oct 22, 2024 · e78ffce · e78ffce
1 parent fba462f
commit e78ffce
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py
@@ -388,7 +388,7 @@ async def set_model_overrides(self, **kwargs):
             if draft_override_args:
                 kwargs["draft_model"] = {
                     **draft_override_args,
-                    **kwargs.get("draft_model"),
+                    **unwrap(kwargs.get("draft_model"), {}),
                 }
 
             # Merge the override and model kwargs

diff --git a/endpoints/OAI/utils/completion.py b/endpoints/OAI/utils/completion.py
@@ -149,8 +149,11 @@ async def load_inline_model(model_name: str, request: Request):
 
         return
 
-    # Load the model
-    await model.load_model(model_path)
+    # Load the model and also add draft dir
+    await model.load_model(
+        model_path,
+        draft_model=config.draft_model.model_dump(include={"draft_model_dir"}),
+    )
 
 
 async def stream_generate_completion(