fix save method of adapter_model.bin

src/llmtuner/tuner/sft/workflow.py

@@ -9,7 +9,8 @@ from llmtuner.extras.misc import get_logits_processor
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.tuner.core import load_model_and_tokenizer
 from llmtuner.tuner.sft.metric import ComputeMetrics
-from llmtuner.tuner.sft.trainer import Seq2SeqPeftTrainer
+from llmtuner.extras.callbacks import SavePeftModelCallback
+from llmtuner.tuner.sft.trainer import CustomSeq2SeqTrainer
 
 if TYPE_CHECKING:
     from transformers import TrainerCallback
@@ -27,6 +28,10 @@ def run_sft(
     dataset = get_dataset(model_args, data_args)
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="sft")
     dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="sft")
+
+    if training_args.predict_with_generate:
+        tokenizer.padding_side = "left" # use left-padding in generation
+
     data_collator = DataCollatorForSeq2Seq(
         tokenizer=tokenizer,
         label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
@@ -41,8 +46,7 @@ def run_sft(
     training_args = Seq2SeqTrainingArguments(**training_args_dict)
 
     # Initialize our Trainer
-    trainer = Seq2SeqPeftTrainer(
-        finetuning_args=finetuning_args,
+    trainer = CustomSeq2SeqTrainer(
         model=model,
         args=training_args,
         tokenizer=tokenizer,
@@ -54,7 +58,7 @@ def run_sft(
 
     # Keyword arguments for `model.generate`
     gen_kwargs = generating_args.to_dict()
-    gen_kwargs["eos_token_id"] = list(set([tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids))
+    gen_kwargs["eos_token_id"] = [tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids
     gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
     gen_kwargs["logits_processor"] = get_logits_processor()
 

src/llmtuner/webui/common.py

@@ -6,7 +6,7 @@ import gradio as gr
 from peft.utils import WEIGHTS_NAME as PEFT_WEIGHTS_NAME
 from transformers.trainer import WEIGHTS_NAME, WEIGHTS_INDEX_NAME
 
-from llmtuner.extras.constants import DEFAULT_TEMPLATE, SUPPORTED_MODELS, DATASET_STAGE_MAP
+from llmtuner.extras.constants import DEFAULT_TEMPLATE, SUPPORTED_MODELS, TRAINING_STAGES
 
 
 DEFAULT_CACHE_DIR = "cache"
@@ -16,8 +16,8 @@ USER_CONFIG = "user.config"
 DATA_CONFIG = "dataset_info.json"
 
 
-def get_save_dir(model_name: str) -> str:
-    return os.path.join(DEFAULT_SAVE_DIR, os.path.split(model_name)[-1])
+def get_save_dir(*args) -> os.PathLike:
+    return os.path.join(DEFAULT_SAVE_DIR, *args)
 
 
 def get_config_path() -> os.PathLike:
@@ -29,7 +29,7 @@ def load_config() -> Dict[str, Any]:
         with open(get_config_path(), "r", encoding="utf-8") as f:
             return json.load(f)
     except:
-        return {"lang": "", "last_model": "", "path_dict": {}}
+        return {"lang": None, "last_model": None, "path_dict": {}, "cache_dir": None}
 
 
 def save_config(lang: str, model_name: str, model_path: str) -> None:
@@ -56,7 +56,7 @@ def get_template(model_name: str) -> str:
 
 def list_checkpoint(model_name: str, finetuning_type: str) -> Dict[str, Any]:
     checkpoints = []
-    save_dir = os.path.join(get_save_dir(model_name), finetuning_type)
+    save_dir = get_save_dir(model_name, finetuning_type)
     if save_dir and os.path.isdir(save_dir):
         for checkpoint in os.listdir(save_dir):
             if (
@@ -78,11 +78,10 @@ def load_dataset_info(dataset_dir: str) -> Dict[str, Any]:
         return {}
 
 
-def list_dataset(dataset_dir: Optional[str] = None, stage: Optional[str] = None) -> Dict[str, Any]:
+def list_dataset(
+    dataset_dir: Optional[str] = None, training_stage: Optional[str] = list(TRAINING_STAGES.keys())[0]
+) -> Dict[str, Any]:
     dataset_info = load_dataset_info(dataset_dir if dataset_dir is not None else DEFAULT_DATA_DIR)
-    if stage:
-        dataset_stage = DATASET_STAGE_MAP[stage]
-        dataset_info = {key: value for key, value in dataset_info.items()
-                        if ("stage" not in value) or value["stage"] == dataset_stage}
-
-    return gr.update(value=[], choices=list(dataset_info.keys()))
\ No newline at end of file
+    ranking = TRAINING_STAGES[training_stage] in ["rm", "dpo"]
+    datasets = [k for k, v in dataset_info.items() if v.get("ranking", False) == ranking]
+    return gr.update(value=[], choices=datasets)

src/llmtuner/webui/components/infer.py

@@ -16,7 +16,7 @@ def create_infer_tab(top_elems: Dict[str, "Component"]) -> Dict[str, "Component"
 
     info_box = gr.Textbox(show_label=False, interactive=False)
 
-    chat_model = WebChatModel()
+    chat_model = WebChatModel(lazy_init=True)
     chat_box, chatbot, history, chat_elems = create_chat_box(chat_model)
 
     load_btn.click(

src/llmtuner/webui/components/train.py

@@ -3,7 +3,7 @@ from transformers.trainer_utils import SchedulerType
 
 import gradio as gr
 
-from llmtuner.extras.constants import STAGES
+from llmtuner.extras.constants import TRAINING_STAGES
 from llmtuner.webui.common import list_checkpoint, list_dataset, DEFAULT_DATA_DIR
 from llmtuner.webui.components.data import create_preview_box
 from llmtuner.webui.utils import can_preview, get_preview, gen_plot
@@ -15,7 +15,9 @@ if TYPE_CHECKING:
 
 def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[str, "Component"]:
     with gr.Row():
-        training_stage = gr.Dropdown(choices=STAGES, value=STAGES[0], scale=2)
+        training_stage = gr.Dropdown(
+            choices=list(TRAINING_STAGES.keys()), value=list(TRAINING_STAGES.keys())[0], scale=2
+        )
         dataset_dir = gr.Textbox(value=DEFAULT_DATA_DIR, scale=2)
         dataset = gr.Dropdown(multiselect=True, scale=4)
         data_preview_btn = gr.Button(interactive=False, scale=1)
@@ -54,7 +56,6 @@ def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dic
             save_steps = gr.Slider(value=100, minimum=10, maximum=5000, step=10)
             warmup_steps = gr.Slider(value=0, minimum=0, maximum=5000, step=1)
             compute_type = gr.Radio(choices=["fp16", "bf16"], value="fp16")
-            padding_side = gr.Radio(choices=["left", "right"], value="left")
 
     with gr.Accordion(label="LoRA config", open=False) as lora_tab:
         with gr.Row():
@@ -120,7 +121,6 @@ def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dic
         save_steps,
         warmup_steps,
         compute_type,
-        padding_side,
         lora_rank,
         lora_dropout,
         lora_target,
@@ -166,7 +166,6 @@ def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dic
         save_steps=save_steps,
         warmup_steps=warmup_steps,
         compute_type=compute_type,
-        padding_side=padding_side,
         lora_tab=lora_tab,
         lora_rank=lora_rank,
         lora_dropout=lora_dropout,