diff --git a/src/llmtuner/data/template.py b/src/llmtuner/data/template.py index 2e1f31c7..ac275c50 100644 --- a/src/llmtuner/data/template.py +++ b/src/llmtuner/data/template.py @@ -682,6 +682,22 @@ register_template( ) +register_template( + name="yi", + prefix=[ + "{{system}}" + ], + prompt=[ + "<|im_start|>user\n{{query}}<|im_end|>\n<|im_start|>assistant\n" + ], + system="", + sep=[ + "<|im_end|>\n" + ], + efficient_eos=True +) + + register_template( name="zephyr", prefix=[ diff --git a/src/llmtuner/extras/constants.py b/src/llmtuner/extras/constants.py index cc9ac290..81040d64 100644 --- a/src/llmtuner/extras/constants.py +++ b/src/llmtuner/extras/constants.py @@ -264,8 +264,12 @@ register_model_group( register_model_group( models={ "Yi-6B": "01-ai/Yi-6B", - "Yi-34B": "01-ai/Yi-34B" - } + "Yi-34B": "01-ai/Yi-34B", + "Yi-34B-Chat": "01-ai/Yi-34B-Chat", + "Yi-34B-int8-Chat": "01-ai/Yi-34B-Chat-8bits", + "Yi-34B-int4-Chat": "01-ai/Yi-34B-Chat-4bits" + }, + template="yi" )