Update app.py

app.py

@@ -9,6 +9,7 @@ TEMPERATURE = 0.5
 TOP_P = 0.95
 TOP_K = 50
 REPETITION_PENALTY = 1.05
+SPECIAL_TOKEN = "->:"
 
 HF_TOKEN = os.getenv('HF_TOKEN')
 
@@ -38,12 +39,12 @@ tokenizer = None
 
 def generate_response(input_text, max_tokens, temperature, top_p, repetition_penalty):
     global model, tokenizer
-    
+
     if model is None or tokenizer is None:
         model, tokenizer = load_model()
-    
-    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
-    
+
+    inputs = tokenizer(input_text + SPECIAL_TOKEN, return_tensors="pt").to(model.device)
+
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -56,12 +57,12 @@ def generate_response(input_text, max_tokens, temperature, top_p, repetition_pen
         )
 
     full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    if "->" in full_response:
-        response_parts = full_response.split("->", 1)
+
+    if SPECIAL_TOKEN in full_response:
+        response_parts = full_response.split(SPECIAL_TOKEN, 1)
         if len(response_parts) > 1:
             return response_parts[1].strip()
-    
+
     return full_response.strip()
 
 def chat_interface(message, history, system_message, max_tokens, temperature, top_p, repetition_penalty):