diff --git a/src/axolotl/cli/preprocess.py b/src/axolotl/cli/preprocess.py index 2c27095191..61925e251e 100644 --- a/src/axolotl/cli/preprocess.py +++ b/src/axolotl/cli/preprocess.py @@ -17,6 +17,7 @@ ) from axolotl.common.cli import PreprocessCliArgs from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH +from axolotl.prompt_strategies.sharegpt import register_chatml_template LOG = logging.getLogger("axolotl.cli.preprocess") @@ -32,6 +33,12 @@ def do_cli(config: Path = Path("examples/"), **kwargs): return_remaining_strings=True ) + if parsed_cfg.chat_template == "chatml" and parsed_cfg.default_system_message: + LOG.info( + f"ChatML set. Adding default system message: {parsed_cfg.default_system_message}" + ) + register_chatml_template(parsed_cfg.default_system_message) + if not parsed_cfg.dataset_prepared_path: msg = ( Fore.RED diff --git a/src/axolotl/cli/train.py b/src/axolotl/cli/train.py index e099a1a6da..df7d6b19eb 100644 --- a/src/axolotl/cli/train.py +++ b/src/axolotl/cli/train.py @@ -16,6 +16,7 @@ print_axolotl_text_art, ) from axolotl.common.cli import TrainerCliArgs +from axolotl.prompt_strategies.sharegpt import register_chatml_template from axolotl.train import train LOG = logging.getLogger("axolotl.cli.train") @@ -32,6 +33,12 @@ def do_cli(config: Path = Path("examples/"), **kwargs): return_remaining_strings=True ) + if parsed_cfg.chat_template == "chatml" and parsed_cfg.default_system_message: + LOG.info( + f"ChatML set. Adding default system message: {parsed_cfg.default_system_message}" + ) + register_chatml_template(parsed_cfg.default_system_message) + if parsed_cfg.rl: dataset_meta = load_rl_datasets(cfg=parsed_cfg, cli_args=parsed_cli_args) else: diff --git a/src/axolotl/prompt_strategies/sharegpt.py b/src/axolotl/prompt_strategies/sharegpt.py index c026889682..15bfee8c47 100644 --- a/src/axolotl/prompt_strategies/sharegpt.py +++ b/src/axolotl/prompt_strategies/sharegpt.py @@ -6,16 +6,19 @@ from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy from axolotl.prompters import ShareGPTPrompterV2 -register_conv_template( - Conversation( - name="chatml", - system_template="<|im_start|>system\n{system_message}", - system_message="You are a helpful assistant.", - roles=["<|im_start|>user", "<|im_start|>assistant"], - sep_style=SeparatorStyle.CHATML, - sep="<|im_end|>", + +def register_chatml_template(system_message=None): + system_message = system_message or "You are a helpful assistant." + register_conv_template( + Conversation( + name="chatml", + system_template="<|im_start|>system\n{system_message}", + system_message=system_message, + roles=["<|im_start|>user", "<|im_start|>assistant"], + sep_style=SeparatorStyle.CHATML, + sep="<|im_end|>", + ) ) -) def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):