Update app.py

app.py

@@ -30,8 +30,8 @@ for model_name, model_endpoint in API_ENDPOINTS.items():
 def format(instruction, history, input, preoutput):
     sy_l, sy_r = SPECIAL_SYMBOLS[0], SPECIAL_SYMBOLS[1]
     formatted_history = "".join(f"{sy_l}{message[0]}{sy_r}\n{sy_l}{message[1]}{sy_r}\n" for message in history)
-    formatted_input = f"{sy_l}System: {instruction}{sy_r}\n{formatted_history}{sy_l}{input}{sy_r}\n{sy_l}{preoutput}"
-    return formatted_input
+    formatted_input = f"{sy_l}System: {instruction}{sy_r}\n{formatted_history}{sy_l}{input}{sy_r}\n"
+    return formatted_input, formatted_input_base = f"{formatted_input}{sy_l}{preoutput}"
     
 def predict(instruction, history, input, preoutput, access_key, model, temperature, top_p, top_k, rep_p, max_tokens, stop_seqs, seed):
 
@@ -46,7 +46,7 @@ def predict(instruction, history, input, preoutput, access_key, model, temperatu
         
     stops = json.loads(stop_seqs)
 
-    formatted_input = format(instruction, history, input, preoutput)
+    formatted_input, formatted_input_base = format(instruction, history, input, preoutput)
     
     response = CLIENTS[model].text_generation(
         formatted_input,
@@ -64,17 +64,19 @@ def predict(instruction, history, input, preoutput, access_key, model, temperatu
     )
 
     sy_l, sy_r = SPECIAL_SYMBOLS[0], SPECIAL_SYMBOLS[1]
-    pre_result = f"{sy_l}{response}{sy_r}{''.join(SPECIAL_SYMBOLS)}"
-    pattern = re.compile(f"{sy_l}(.*){sy_r}", re.DOTALL)
-    match = pattern.search(pre_result)
-    get_result = preoutput + match.group(1).strip() if match else ""
-
-    history = history + [[input, get_result]]
+    result = ""
+    
+    for stop in stops:
+        result = response.split(stop, 1)[0]
+    for symbol in stops:
+        result = response.replace(symbol, '')
+    
+    history = history + [[input, result]]
 
-    print(formatted_input + get_result)
-    print(f"---\nUSER: {input}\nBOT: {get_result}\n---")
+    print(formatted_input_base + result)
+    print(f"---\nUSER: {input}\nBOT: {result}\n---")
 
-    return (get_result, input, history)
+    return (result, input, history)
 
 def clear_history():
     print(">>> HISTORY CLEARED!")