Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import re
 
 MODEL = "prithivMLmods/Llama-SmolTalk-3.2-1B-Instruct"
 
@@ -18,37 +19,37 @@ chatbot = pipeline(
     device_map="auto"
 )
 
-system_prompt = """Tu es Aria, une IA gentille, claire et polie. 
-Réponds toujours en phrases complètes. 
-Ne fais pas de jeu de rôle. 
-Ne répète pas le message de l'utilisateur. 
-Réponds directement à la question."""
+system_prompt = """Tu es Aria, une IA bienveillante, claire et polie.
+Réponds directement aux questions en une ou plusieurs phrases complètes.
+Ne répète jamais le message de l'utilisateur. Ne joue pas un rôle.
+N'inclus pas "Utilisateur:" ou "Aria:" dans ta réponse."""
+
+def clean_reply(text):
+    # Supprimer tout ce qui pourrait être un rôle ou répétition
+    text = re.sub(r"^.*?Aria:\s*", "", text, flags=re.DOTALL)
+    text = re.sub(r"Utilisateur:.*", "", text, flags=re.DOTALL)
+    text = text.strip()
+    return text
 
 def chat(message, history=[]):
-    history = history[-3:]  # garder 3 derniers échanges
-    
-    context = "\n".join([f"Utilisateur: {m[0]}\nAria: {m[1]}" for m in history])
-    prompt = f"{system_prompt}\n{context}\nUtilisateur: {message}\nAria:"
+    history = history[-3:]
+
+    # Contexte plus naturel
+    context = "\n".join([f"{m[0]}\nRéponse: {m[1]}" for m in history])
+    prompt = f"{system_prompt}\n\n{context}\n{message}\nRéponse:"
 
     resp = chatbot(
         prompt,
-        max_new_tokens=120,  # plus long
+        max_new_tokens=120,
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
         repetition_penalty=1.1
     )[0]["generated_text"]
 
-    # Garde seulement après "Aria:"
-    reply = resp.split("Aria:")[-1].strip()
-
-    # Supprime si "Utilisateur:" est revenu dans la génération
-    if "Utilisateur:" in reply:
-        reply = reply.split("Utilisateur:")[0].strip()
-
+    reply = clean_reply(resp)
     history.append([message, reply])
 
-    # Retourner au format Gradio Chatbot
     return history, history
 
 with gr.Blocks() as demo: