From 754fb15f234b5d725f825af8594bb9b344aa0be6 Mon Sep 17 00:00:00 2001 From: kingbri Date: Tue, 17 Sep 2024 22:48:56 -0400 Subject: [PATCH] Config: Fix draft model migration and loading The loader takes in the "draft" parameter, so map the config model to that when creating kwargs for initial load. Also map the old "draft" key to the new "draft_model" key. Signed-off-by: kingbri --- common/tabby_config.py | 2 +- main.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/common/tabby_config.py b/common/tabby_config.py index 0bf3563..43f227e 100644 --- a/common/tabby_config.py +++ b/common/tabby_config.py @@ -72,7 +72,7 @@ def _from_file(self, config_path: pathlib.Path): if model_cfg.get("draft"): legacy = True - cfg["draft"] = model_cfg["draft"] + cfg["draft_model"] = model_cfg["draft"] if model_cfg.get("lora"): legacy = True diff --git a/main.py b/main.py index e340b70..b83fda7 100644 --- a/main.py +++ b/main.py @@ -69,7 +69,11 @@ async def entrypoint_async(): model_path = model_path / model_name # TODO: remove model_dump() - await model.load_model(model_path.resolve(), **config.model.model_dump()) + await model.load_model( + model_path.resolve(), + **config.model.model_dump(), + draft=config.draft_model.model_dump(), + ) # Load loras after loading the model if config.lora.loras: