Update app.py

app.py

@@ -81,22 +81,29 @@ def generate_response(
     # Remove leftover special tokens
     response = response.replace("<|im_end|>", "").strip()
 
-    # Remove any label and text following it on the same line
     lines = response.splitlines()
+
+    # If the first line starts with Assistant:, return only the rest of that line
+    first_line = lines[0].strip() if lines else ""
+    for label in ["Assistant:", "assistant:"]:
+        if first_line.lower().startswith(label.lower()):
+            return first_line[len(label):].strip()
+
+    # Otherwise, clean each line by removing any label and following text
     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
+        if line:
             cleaned_lines.append(line)
 
     response = "\n".join(cleaned_lines)
-
     return response
 
 
 
+
 # Respond function for Gradio
 def respond(
     message,