Update app.py

app.py

@@ -19,7 +19,22 @@ generator = pipeline(
     tokenizer=tokenizer,
     device=device
 )
-
+def clean_reply(text):
+    
+    text = text.strip()
+    
+    for prefix in ["Assistant:", "assistant:", "User:", "user:"]:
+        if text.startswith(prefix):
+            text = text[len(prefix):].strip()
+            
+    lines = [l.strip() for l in text.split("\n")]
+    lines = [l for l in lines if l]
+
+    if len(lines) == 0:
+        return ""
+
+    return lines[0] 
+    
 def chat_with_model(user_message, chat_history, max_new_tokens=60, temperature=0.8, top_p=0.9):
 
     if chat_history is None: