Update app.py

app.py

@@ -3,17 +3,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 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,
@@ -21,41 +18,42 @@ chatbot = pipeline(
     device_map="auto"
 )
 
-# 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."
-)
+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."""
 
-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:"
+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:"
 
     resp = chatbot(
         prompt,
-        max_new_tokens=250,   # plus long pour éviter les coupures
+        max_new_tokens=120,  # plus long
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
         repetition_penalty=1.1
     )[0]["generated_text"]
 
-    # Couper dès qu'il repart sur un nouveau tour
-    reply = resp.split("Réponse:")[-1].strip()
-    reply = reply.split("Utilisateur:")[0].strip()
+    # 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()
 
-    history.append((message, reply))
+    history.append([message, reply])
+
+    # Retourner au format Gradio Chatbot
     return history, history
 
 with gr.Blocks() as demo:
-    chatbot_ui = gr.Chatbot()
-    state = gr.State([])  # sauvegarde de l'historique
+    chat_ui = gr.Chatbot()
     msg = gr.Textbox(placeholder="Écris un message...")
-
-    msg.submit(chat, [msg, state], [chatbot_ui, state])
-    msg.submit(lambda: "", None, msg)  # reset input après envoi
+    msg.submit(chat, [msg, chat_ui], [chat_ui, chat_ui])
 
 demo.launch()