Update app.py

app.py

@@ -1,17 +1,19 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import re
 
 MODEL = "prithivMLmods/Llama-SmolTalk-3.2-1B-Instruct"
 
+# Charger le tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 
+# Charger le modèle en 8 bits
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     device_map="auto",
     load_in_8bit=True
 )
 
+# Pipeline
 chatbot = pipeline(
     "text-generation",
     model=model,
@@ -19,33 +21,41 @@ chatbot = pipeline(
     device_map="auto"
 )
 
-system_prompt = """Tu es Aria, une IA gentille et claire. 
-Réponds directement à la question de l'utilisateur en français, 
-en une ou deux phrases complètes. 
-Ne répète pas la question. Ne joue pas de rôle."""
-
-def clean_reply(text):
-    # Retirer tout texte avant ou après la vraie réponse
-    text = re.sub(r"^.*?:", "", text, flags=re.DOTALL)  # supprime jusqu'au premier ":"
-    return text.strip()
+# Prompt système
+system_prompt = (
+    "Tu es Aria, une IA gentille, claire et polie. "
+    "Réponds toujours en phrases complètes. "
+    "Ne te lances pas dans un jeu de rôle, ne répète pas les messages précédents, "
+    "et donne uniquement ta réponse."
+)
 
-def chat(message):
-    prompt = f"{system_prompt}\n\nQuestion : {message}\nRéponse :"
+def chat(message, history):
+    history = history or []
+    context = "\n".join([f"{user}\n{bot}" for user, bot in history[-3:]])
+    prompt = f"{system_prompt}\n{context}\n{message}\nRéponse:"
 
     resp = chatbot(
         prompt,
-        max_new_tokens=80,
+        max_new_tokens=250,   # plus long pour éviter les coupures
         do_sample=True,
         temperature=0.7,
-        top_p=0.9
+        top_p=0.9,
+        repetition_penalty=1.1
     )[0]["generated_text"]
 
-    reply = clean_reply(resp)
-    return reply
+    # Couper dès qu'il repart sur un nouveau tour
+    reply = resp.split("Réponse:")[-1].strip()
+    reply = reply.split("Utilisateur:")[0].strip()
+
+    history.append((message, reply))
+    return history, history
 
 with gr.Blocks() as demo:
+    chatbot_ui = gr.Chatbot()
+    state = gr.State([])  # sauvegarde de l'historique
     msg = gr.Textbox(placeholder="Écris un message...")
-    output = gr.Textbox(label="Réponse d'Aria")
-    msg.submit(lambda m: chat(m), msg, output)
+
+    msg.submit(chat, [msg, state], [chatbot_ui, state])
+    msg.submit(lambda: "", None, msg)  # reset input après envoi
 
 demo.launch()