fix webui

Former-commit-id: e87630ef77977b2879f1199b9a421acbbbb32a51
2023-08-03 12:43:12 +08:00
parent 28a51b622b
commit 9c96b97dc7
3 changed files with 3 additions and 2 deletions
--- a/src/llmtuner/tuner/core/loader.py
+++ b/src/llmtuner/tuner/core/loader.py
@@ -128,7 +128,7 @@ def load_model_and_tokenizer(
    model = init_adapter(model, model_args, finetuning_args, is_trainable, is_mergeable)

    if stage == "rm" or stage == "ppo": # add value head
-        model = AutoModelForCausalLMWithValueHead.from_pretrained(model)
+        model: "AutoModelForCausalLMWithValueHead" = AutoModelForCausalLMWithValueHead.from_pretrained(model)
        reset_logging()

        if stage == "rm" and model_args.checkpoint_dir is not None: # load valuehead weights to evaluate reward model