From 2310b3bc31646ea7134cf8c3c12a73587539cd79 Mon Sep 17 00:00:00 2001
From: bys0318 <bys0318@126.com>
Date: Thu, 2 Nov 2023 15:27:39 +0800
Subject: [PATCH] support chatglm3
---
pred.py | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/pred.py b/pred.py
index 467d38b..8163891 100644
--- a/pred.py
+++ b/pred.py
@@ -17,7 +17,9 @@ def parse_args(args=None):
# This is the customized building prompt for chat models
def build_chat(tokenizer, prompt, model_name):
- if "chatglm" in model_name:
+ if "chatglm3" in model_name:
+ prompt = tokenizer.build_chat_input(prompt)
+ elif "chatglm" in model_name:
prompt = tokenizer.build_prompt(prompt)
elif "longchat" in model_name or "vicuna" in model_name:
from fastchat.model import get_conversation_template
--
GitLab