diff --git a/lmdeploy/model.py b/lmdeploy/model.py index 626f3f5b1b..c1700b37d6 100644 --- a/lmdeploy/model.py +++ b/lmdeploy/model.py @@ -1583,10 +1583,12 @@ def __init__(self, system='<|system|>\n', user='<|user|>\n', assistant='<|assistant|>\n', + stop_words=['<|user|>', '<|endoftext|>', '<|observation|>'], **kwargs): super().__init__(system=system, user=user, assistant=assistant, + stop_words=stop_words, **kwargs) self.start = '[gMASK]' diff --git a/tests/test_lmdeploy/test_model.py b/tests/test_lmdeploy/test_model.py index f38b952e33..3a00266d19 100644 --- a/tests/test_lmdeploy/test_model.py +++ b/tests/test_lmdeploy/test_model.py @@ -329,7 +329,13 @@ def test_chatglm3(): def test_glm4(): - model = MODELS.get('glm4')() + model_path_and_name = 'THUDM/glm-4-9b-chat' + deduced_name = best_match_model(model_path_and_name) + assert deduced_name == 'glm4' + + model = MODELS.get(deduced_name)() + # check stop words + assert model.stop_words == ['<|user|>', '<|endoftext|>', '<|observation|>'] messages = [{ 'role': 'system', 'content': 'you are a helpful assistant' @@ -344,7 +350,7 @@ def test_glm4(): 'content': 'AGI is?' }] from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained('THUDM/glm-4-9b-chat', + tokenizer = AutoTokenizer.from_pretrained(model_path_and_name, trust_remote_code=True) ref = tokenizer.apply_chat_template(messages, tokenize=False) res = model.messages2prompt(messages)