diff --git a/__init__.py b/__init__.py index 081e814..4315115 100644 --- a/__init__.py +++ b/__init__.py @@ -83,6 +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": ""}), }, } diff --git a/mz_llama_cpp.py b/mz_llama_cpp.py index 5840895..35e05c7 100644 --- a/mz_llama_cpp.py +++ b/mz_llama_cpp.py @@ -482,11 +482,7 @@ def base_query_beautify_prompt_text(args_dict): mz_prompt_utils.Utils.print_log(f"question: {question}") - - - if schema is not None: - response_json = llama_cpp_simple_interrogator_to_json( model_file=model_file, system=system_prompt, @@ -494,9 +490,7 @@ def base_query_beautify_prompt_text(args_dict): schema=schema, options=options, ) - mz_prompt_utils.Utils.print_log(f"response_json: {response_json}") - if keep_device is False: - freed_gpu_memory(model_file=model_file) + mz_prompt_utils.Utils.print_log(f"response_json: {response_json}") response = json.loads(response_json) full_responses = [] @@ -535,8 +529,20 @@ def base_query_beautify_prompt_text(args_dict): options=options, ) + prefix = customize_instruct.get("prefix", "") + # 去除前缀 + if prefix != "": + full_response = full_response.lstrip(prefix) + + suffix = customize_instruct.get("suffix", "") + # 去除后缀 + if suffix != "": + full_response = full_response.rstrip(suffix) + + + if keep_device is False: + freed_gpu_memory(model_file=model_file) - # 去除换行 while full_response.find("\n") != -1: full_response = full_response.replace("\n", " ")