Merge pull request #221 from xming521/dev

Dev

.gitignore

@@ -197,3 +197,4 @@ models_final/*
 /data/*
 /llamaboard_cache
 eval_Result/*
+.claude/*

weclone/cli.py

@@ -120,6 +120,15 @@ def train_sft():
     train_sft_main()
 
 
+@cli.command("train-pt", help="Continue pre-training the model using prepared text datasets.")
+@apply_common_decorators()
+def train_pt():
+    """Continue pre-training the model using prepared text datasets."""
+    from weclone.train.train_pt import main as train_pt_main
+
+    train_pt_main()
+
+
 @cli.command("webchat-demo", help="Launch Web UI for interactive testing with fine-tuned model.")
 @apply_common_decorators()
 def web_demo():

weclone/train/train_pt.py

@@ -0,0 +1,72 @@
+import json
+from pathlib import Path
+from typing import Any, cast
+
+from llamafactory.extras.misc import get_current_device
+from llamafactory.train.tuner import run_exp
+
+from weclone.utils.config import load_config
+from weclone.utils.config_models import WCTrainPtConfig
+from weclone.utils.log import logger
+
+
+def _resolve_dataset_path(dataset_dir: str, dataset_name: str) -> Path:
+    dataset_info_path = Path(dataset_dir) / "dataset_info.json"
+    if not dataset_info_path.exists():
+        raise FileNotFoundError(f"Dataset info file does not exist: {dataset_info_path}")
+
+    with dataset_info_path.open("r", encoding="utf-8") as f:
+        dataset_info: dict[str, Any] = json.load(f)
+
+    dataset_entry = dataset_info.get(dataset_name)
+    if dataset_entry is None:
+        raise ValueError(f"Dataset '{dataset_name}' is not defined in {dataset_info_path}")
+
+    if dataset_entry.get("formatting") == "sharegpt":
+        raise ValueError(
+            f"Dataset '{dataset_name}' is a ShareGPT dataset. "
+            "LlamaFactory pre-training requires Alpaca-style data with columns.prompt mapped to text."
+        )
+
+    prompt_column = (dataset_entry.get("columns") or {}).get("prompt")
+    if prompt_column is None:
+        raise ValueError(f"Dataset '{dataset_name}' must define columns.prompt for pre-training.")
+
+    dataset_file_name = dataset_entry.get("file_name")
+    if not dataset_file_name:
+        raise ValueError(f"Dataset '{dataset_name}' must define file_name in {dataset_info_path}")
+
+    data_path = Path(dataset_file_name)
+    if not data_path.is_absolute():
+        data_path = Path(dataset_dir) / data_path
+
+    if not data_path.exists():
+        raise FileNotFoundError(f"Dataset file '{data_path}' does not exist.")
+
+    return data_path
+
+
+def main():
+    train_config = cast(WCTrainPtConfig, load_config(arg_type="train_pt"))
+
+    if train_config.stage != "pt":
+        raise ValueError(f"train-pt requires stage='pt', got stage={train_config.stage!r}")
+
+    device = get_current_device()
+    if device == "cpu":
+        logger.warning("Please note you are using CPU for training, non-Mac devices may encounter issues")
+
+    data_path = _resolve_dataset_path(train_config.dataset_dir, train_config.dataset)
+    logger.info(f"Using pre-training dataset: {data_path}")
+
+    formatted_config = json.dumps(train_config.model_dump(mode="json"), indent=4, ensure_ascii=False)
+    logger.info(f"Continued pre-training configuration:\n{formatted_config}")
+
+    config_dict = train_config.model_dump(mode="json", exclude_none=True)
+    config_dict.pop("quantization", None)
+
+    run_exp(config_dict)
+
+
+if __name__ == "__main__":
+    main()

weclone/utils/config.py

@@ -10,6 +10,7 @@ from .config_models import (
     WcConfig,
     WCInferConfig,
     WCMakeDatasetConfig,
+    WCTrainPtConfig,
     WCTrainSftConfig,
 )
 from .log import logger
@@ -93,6 +94,18 @@ def create_config_by_arg_type(arg_type: str, wc_config: WcConfig) -> BaseModel:
         config_dict = {**common_config, **train_dict}
         return WCTrainSftConfig(**config_dict)
 
+    elif arg_type == "train_pt":
+        if wc_config.train_pt_args is None:
+            logger.error("Missing `train_pt_args` in configuration file.")
+            sys.exit(1)
+
+        train_dict = wc_config.train_pt_args.model_dump()
+        train_dict.pop("quantization", None)
+        train_dict.update(_flatten_quantization_args(wc_config.train_pt_args))
+
+        config_dict = {**common_config, **train_dict}
+        return WCTrainPtConfig(**config_dict)
+
     elif arg_type == "make_dataset":
         make_dataset_config = wc_config.make_dataset_args.model_dump()
         train_sft_args = wc_config.train_sft_args

weclone/utils/config_models.py

@@ -202,6 +202,10 @@ class TrainSftArgs(BaseConfigModel):
     stage: str = Field("sft", description="Training stage")
     dataset: str = Field(..., description="Dataset name")
     dataset_dir: str = Field("./dataset/res_csv/sft", description="Dataset directory")
+    resume_adapter_name_or_path: Optional[str] = Field(
+        None,
+        description="Existing LoRA adapter path to continue SFT from. Output still uses common_args.adapter_name_or_path.",
+    )
     freeze_multi_modal_projector: bool = Field(
         False, description="Whether to freeze multimodal projector during MLLM training"
     )
@@ -235,6 +239,16 @@ class TrainSftArgs(BaseConfigModel):
     do_train: bool = Field(True)
 
 
+class TrainPtArgs(TrainSftArgs):
+    stage: str = Field("pt", description="Pre-training stage")
+    dataset: str = Field(..., description="Pre-training dataset name")
+    output_dir: Optional[str] = Field(None, description="PT output directory")
+    packing: Optional[bool] = Field(
+        None,
+        description="Whether to pack sequences. LlamaFactory enables packing automatically for stage=pt.",
+    )
+
+
 class InferArgs(BaseConfigModel):
     repetition_penalty: float = Field(1.2, description="Repetition penalty")
     temperature: float = Field(..., description="Temperature")
@@ -272,6 +286,7 @@ class WcConfig(BaseModel):
     cli_args: CliArgs = Field(..., description="Command line arguments")
     make_dataset_args: MakeDatasetArgs = Field(..., description="Dataset processing parameters")
     train_sft_args: TrainSftArgs = Field(..., description="SFT fine-tuning parameters")
+    train_pt_args: Optional[TrainPtArgs] = Field(None, description="PT continued pre-training parameters")
     infer_args: InferArgs = Field(..., description="Inference parameters")
     vllm_args: VllmArgs = Field(VllmArgs())
     test_model_args: TestModelArgs = Field(TestModelArgs())
@@ -292,21 +307,48 @@ class WCTrainSftConfig(CommonArgs, TrainSftArgs, CommonMethods):
 
     @model_validator(mode="after")
     def process_config(self):
-        adapter_name_value = getattr(self, "adapter_name_or_path", None)
+        output_adapter_value = getattr(self, "adapter_name_or_path", None)
+        resume_adapter_value = getattr(self, "resume_adapter_name_or_path", None)
 
-        if adapter_name_value:
-            self.output_dir = adapter_name_value
+        if output_adapter_value:
+            self.output_dir = output_adapter_value
+
+        if resume_adapter_value:
+            self.adapter_name_or_path = resume_adapter_value
+        elif hasattr(self, "adapter_name_or_path"):
+            delattr(self, "adapter_name_or_path")
 
         self.dataset = self._parse_dataset_name()
-        # Always remove adapter_name_or_path field after processing
-        if hasattr(self, "adapter_name_or_path"):
-            delattr(self, "adapter_name_or_path")
+        if hasattr(self, "resume_adapter_name_or_path"):
+            delattr(self, "resume_adapter_name_or_path")
+        if hasattr(self, "quantization"):
+            delattr(self, "quantization")
         if hasattr(self, "include_type"):
             delattr(self, "include_type")
 
         return self
 
 
+class WCTrainPtConfig(CommonArgs, TrainPtArgs):
+    """Final configuration model for continued pre-training"""
+
+    output_dir: Optional[str] = Field(None)
+
+    @model_validator(mode="after")
+    def process_config(self):
+        adapter_name_value = getattr(self, "adapter_name_or_path", None)
+
+        if self.output_dir is None and adapter_name_value:
+            self.output_dir = adapter_name_value
+
+        if hasattr(self, "adapter_name_or_path"):
+            delattr(self, "adapter_name_or_path")
+        if hasattr(self, "quantization"):
+            delattr(self, "quantization")
+
+        return self
+
+
 class WCMakeDatasetConfig(CommonArgs, MakeDatasetArgs, CommonMethods):
     """Final configuration model for creating datasets"""