diff --git a/.gitignore b/.gitignore index 5722078..fba18e5 100644 --- a/.gitignore +++ b/.gitignore @@ -197,3 +197,4 @@ models_final/* /data/* /llamaboard_cache eval_Result/* +.claude/* diff --git a/WC-exp b/WC-exp index e0acf48..56af4ad 160000 --- a/WC-exp +++ b/WC-exp @@ -1 +1 @@ -Subproject commit e0acf48c642b56625832e858ce3a7bf95ac14270 +Subproject commit 56af4ad68fc570a669c311f64ec27ea9eb9586a6 diff --git a/weclone/cli.py b/weclone/cli.py index eb0f93f..a758a1e 100644 --- a/weclone/cli.py +++ b/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(): diff --git a/weclone/train/train_pt.py b/weclone/train/train_pt.py new file mode 100644 index 0000000..874fdea --- /dev/null +++ b/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() diff --git a/weclone/utils/config.py b/weclone/utils/config.py index 53825a9..d6ee3a6 100644 --- a/weclone/utils/config.py +++ b/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 diff --git a/weclone/utils/config_models.py b/weclone/utils/config_models.py index 6d7ef0a..137a560 100644 --- a/weclone/utils/config_models.py +++ b/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"""