Upload 5 files

app.py

@@ -1201,33 +1201,42 @@ def format_message(messages):
 
 def extract_role(messages):
     role_map = {"user": "Human", "assistant": "Assistant", "system": "System"}
-    prefix = False
+    prefix = True  # 默认添加前缀
     first_message = messages[0]["content"]
     pattern = re.compile(
         r"""
         <roleInfo>\s*
-        user:\s*(?P<user>[^\n]*)\s*
-        assistant:\s*(?P<assistant>[^\n]*)\s*
-        system:\s*(?P<system>[^\n]*)\s*
-        prefix:\s*(?P<prefix>[^\n]*)\s*
+        (?:user:\s*(?P<user>[^\n]*)\s*)?        # Make user optional
+        (?:assistant:\s*(?P<assistant>[^\n]*)\s*)? # Make assistant optional
+        (?:system:\s*(?P<system>[^\n]*)\s*)?     # Make system optional
+        (?:prefix:\s*(?P<prefix>[^\n]*)\s*)?     # Make prefix optional
         </roleInfo>\n
     """,
         re.VERBOSE,
     )
     match = pattern.search(first_message)
     if match:
-        role_map = {
-            "user": match.group("user"),
-            "assistant": match.group("assistant"),
-            "system": match.group("system"),
-        }
-        prefix = match.group("prefix") == "1"
+        # 更新 role_map 如果提供了值
+        user_role = match.group("user")
+        assistant_role = match.group("assistant")
+        system_role = match.group("system")
+        if user_role: role_map["user"] = user_role
+        if assistant_role: role_map["assistant"] = assistant_role
+        if system_role: role_map["system"] = system_role
+
+        # 检查 prefix 值：仅当显式设置为非 "1" 时才将 prefix 设为 False
+        prefix_value = match.group("prefix")
+        if prefix_value is not None and prefix_value != "1":
+            prefix = False
+        # 如果 prefix_value 是 None (标签不存在) 或 "1", prefix 保持 True
+
         messages[0]["content"] = pattern.sub("", first_message)
         print(f"Extracted role map:")
         print(
             f"User: {role_map['user']}, Assistant: {role_map['assistant']}, System: {role_map['system']}"
         )
-        print(f"Using prefix: {prefix}")
+        print(f"Using prefix: {prefix}") # 打印语句保持不变，反映最终结果
+    # 如果没有匹配到 <roleInfo>，prefix 保持默认值 True
     return (role_map, prefix, messages)