From 2d19ba04e3dc25f615085b67a53d59b46d30067b Mon Sep 17 00:00:00 2001
From: hoshi-hiyouga <hiyouga@buaa.edu.cn>
Date: Thu, 19 Sep 2024 02:13:39 +0800
Subject: [PATCH] fix webui

Former-commit-id: f0b930d94a31ed46a605b3e8e3d989fdc344e61b
---
 src/llamafactory/webui/common.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/llamafactory/webui/common.py b/src/llamafactory/webui/common.py
index 019812c7..0ad2929e 100644
--- a/src/llamafactory/webui/common.py
+++ b/src/llamafactory/webui/common.py
@@ -137,8 +137,13 @@ def get_template(model_name: str) -> str:
     r"""
     Gets the template name if the model is a chat model.
     """
-    if model_name and model_name.endswith("Chat") and get_prefix(model_name) in DEFAULT_TEMPLATE:
+    if (
+        model_name
+        and any(suffix in model_name for suffix in ("-Chat", "-Instruct"))
+        and get_prefix(model_name) in DEFAULT_TEMPLATE
+    ):
         return DEFAULT_TEMPLATE[get_prefix(model_name)]
+
     return "default"