From 1a14e7de8fc5b08503dd4e10c1f7332a301b911a Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Tue, 30 Apr 2024 23:03:36 +0800
Subject: [PATCH] 'update'

---
 __init__.py     |  4 ++--
 mz_llama_cpp.py | 18 ++++++++++--------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/__init__.py b/__init__.py
index 8d08352..e114c57 100644
--- a/__init__.py
+++ b/__init__.py
@@ -83,8 +83,8 @@ class MZ_CustomizeInstruct:
             "required": {
                 "system": ("STRING", {"multiline": True, "default": mz_prompts.Long_prompt}),
                 "instruct": ("STRING", {"multiline": True, "default": "Short: {text}"}),
-                "prefix": ("STRING", {"default": "Long: "}),
-                "suffix": ("STRING", {"default": ""}),
+                "start_str": ("STRING", {"default": "Long: "}),
+                "end_str": ("STRING", {"default": ""}),
             },
         }
 
diff --git a/mz_llama_cpp.py b/mz_llama_cpp.py
index 35e05c7..638aa6d 100644
--- a/mz_llama_cpp.py
+++ b/mz_llama_cpp.py
@@ -529,15 +529,17 @@ def base_query_beautify_prompt_text(args_dict):
                 options=options,
             )
 
-            prefix = customize_instruct.get("prefix", "")
-            # 去除前缀
-            if prefix != "":
-                full_response = full_response.lstrip(prefix)
+            start_str = customize_instruct.get("start_str", "") 
+            if start_str != "" and full_response.find(start_str) != -1:
+                full_response_list = full_response.split(start_str)
+                full_response_list[0] = ""
+                full_response = start_str.join(full_response_list)
 
-            suffix = customize_instruct.get("suffix", "")
-            # 去除后缀
-            if suffix != "":
-                full_response = full_response.rstrip(suffix)
+            end_str = customize_instruct.get("end_str", "")
+            if end_str != "" and full_response.find(end_str) != -1:
+                full_response_list = full_response.split(end_str)
+                full_response_list[-1] = ""
+                full_response = end_str.join(full_response_list)
 
 
         if keep_device is False: