Update app.py

app.py

@@ -37,11 +37,11 @@ SYSTEM_PROMPT = (
 
 
 def build_context(user_message):
-    return f"{SYSTEM_PROMPT}\n\nUser: {user_message}"
+    return SYSTEM_PROMPT + "\n\nUser: " + user_message + "\n\nAssistant:"
+
 
 
 
-# Generate response
 def generate_response(
     prompt,
     max_tokens=300,
@@ -74,13 +74,27 @@ def generate_response(
             eos_token_id=tokenizer.eos_token_id,
         )
 
+    # Take only newly generated tokens
     new_tokens = outputs[0][inputs.input_ids.shape[-1]:]
     response = tokenizer.decode(new_tokens, skip_special_tokens=True)
-    
- 
-    response = response.replace("<|im_end|>", "")
-    
-    return response.strip()
+
+    # Remove leftover special tokens
+    response = response.replace("<|im_end|>", "").strip()
+
+    # Remove any label and text following it on the same line
+    lines = response.splitlines()
+    cleaned_lines = []
+    for line in lines:
+        for label in ["Assistant:", "assistant:", "User:", "user:"]:
+            if label in line:
+                line = line.split(label)[0].strip()
+        if line:  # keep non-empty lines
+            cleaned_lines.append(line)
+
+    response = "\n".join(cleaned_lines)
+
+    return response
+
 
 
 # Respond function for Gradio