refactor dataset_attr, add eos in pt, fix #757

Former-commit-id: 0feec9a830b917b36686b61938a66e842eccf930
2023-09-01 19:00:45 +08:00
parent 93be211f80
commit e5b72c6a77
19 changed files with 108 additions and 126 deletions
--- a/src/llmtuner/extras/constants.py
+++ b/src/llmtuner/extras/constants.py
@@ -10,20 +10,12 @@ LAYERNORM_NAMES = ["norm", "ln_f", "ln_attn", "ln_mlp"]

 METHODS = ["full", "freeze", "lora"]

-STAGES = [
-    "SFT",
-    "Reward Modeling",
-    "PPO",
-    "DPO",
-    "Pre-Training"
-]
-
-DATASET_STAGE_MAP = {
-    "SFT": "sft",
-    "Pre-Training": "pt",
+TRAINING_STAGES = {
+    "Supervised Fine-Tuning": "sft",
    "Reward Modeling": "rm",
-    "PPO": "sft",
-    "DPO": "rm"
+    "PPO": "ppo",
+    "DPO": "dpo",
+    "Pre-Training": "pt"
 }

 SUPPORTED_MODELS = {