From ae6ad55758e316d0f631aa0267025942919860e9 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: aa6e65b24451fe9f65d58e5eca5a56eb9aba71e8
---
 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"