handle if assistant_message_gen and assistant_message_gen!=assistant_message_start, replace final output tag with unspaced (gen) version if exists

Author: LostRuins

kcpp_adapters/Jamba.json

@@ -0,0 +1,9 @@
+{
+  "system_start": "<|bom|><|system|> ",
+  "system_end": "<|eom|>",
+  "user_start": "<|bom|><|user|> ",
+  "user_end": "<|eom|>",
+  "assistant_start": "<|bom|><|assistant|> ",
+  "assistant_gen": "<|bom|><|assistant|>",
+  "assistant_end": "<|eom|>"
+}

kcpp_adapters/Mistral-NonTekken.json

@@ -3,6 +3,7 @@
   "system_end": "",
   "user_start": "[INST] ",
   "user_end": "",
-  "assistant_start": "[/INST]",
+  "assistant_start": "[/INST] ",
+  "assistant_gen": "[/INST]",
   "assistant_end": "</s>"
 }

kcpp_adapters/RWKV-World.json

@@ -3,6 +3,7 @@
   "system_end": "\n\n",
   "user_start": "User: ",
   "user_end": "\n\n",
-  "assistant_start": "Assistant:",
+  "assistant_start": "Assistant: ",
+  "assistant_gen": "Assistant:",
   "assistant_end": "\n\n"
 }

koboldcpp.py

@@ -671,6 +671,14 @@ def tryparsefloat(value,fallback):
     except ValueError:
         return fallback
 
+def replace_last_in_string(text: str, match: str, replacement: str) -> str:
+    if match == "":
+        return text
+    head, sep, tail = text.rpartition(match)
+    if sep == "":
+        return text  # old not found
+    return head + replacement + tail
+
 def is_incomplete_utf8_sequence(byte_seq): #note, this will only flag INCOMPLETE sequences, corrupted ones will be ignored.
     try:
         byte_seq.decode('utf-8')
@@ -2608,6 +2616,11 @@ def transform_genparams(genparams, api_format):
         assistant_message_start = adapter_obj.get("assistant_start", "\n### Response:\n")
         assistant_message_end = adapter_obj.get("assistant_end", "")
         if isinstance(prompt, str): #needed because comfy SD uses same field name
+            if assistant_message_gen and assistant_message_gen!=assistant_message_start: #replace final output tag with unspaced (gen) version if exists
+                if prompt.rstrip().endswith("{{[OUTPUT]}}"):
+                    prompt = replace_last_in_string(prompt,"{{[OUTPUT]}}",assistant_message_gen)
+                elif assistant_message_start and prompt.rstrip().endswith(assistant_message_start):
+                    prompt = replace_last_in_string(prompt, assistant_message_start, assistant_message_gen)
             if "{{[INPUT_END]}}" in prompt or "{{[OUTPUT_END]}}" in prompt:
                 prompt = prompt.replace("{{[INPUT]}}", user_message_start)
                 prompt = prompt.replace("{{[OUTPUT]}}", assistant_message_start)