Update llm_utils.py

llm_toolkit/llm_utils.py

@@ -10,6 +10,17 @@ from transformers import (
 from tqdm import tqdm
 
 
+def get_template(model_name):
+    model_name = model_name.lower()
+    if "llama-3" in model_name:
+        return "llama3"
+    if "internlm" in model_name:
+        return "intern2"
+    if "glm" in model_name:
+        return "glm4"
+    return "chatml"
+
+
 def load_model(
     model_name,
     dtype=torch.bfloat16,
@@ -22,7 +33,7 @@ def load_model(
     if adapter_name_or_path and using_llama_factory:
         from llamafactory.chat import ChatModel
 
-        template = "llama3" if "llama-3" in model_name.lower() else "chatml"
+        template = get_template(model_name)
 
         args = dict(
             model_name_or_path=model_name,