Update app.py

app.py

@@ -6,14 +6,14 @@ MODEL = "prithivMLmods/Llama-SmolTalk-3.2-1B-Instruct"
 # Charger le tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 
-# Charger le modèle en 8 bits pour accélérer et réduire l’usage mémoire
+# Charger le modèle en 8 bits
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     device_map="auto",
-    load_in_8bit=True  # accélère sur CPU/peu de mémoire
+    load_in_8bit=True
 )
 
-# Pipeline de génération
+# Pipeline
 chatbot = pipeline(
     "text-generation",
     model=model,
@@ -22,34 +22,40 @@ chatbot = pipeline(
 )
 
 # Prompt système
-system_prompt = """Tu es Aria, une IA gentille, claire et polie. Réponds toujours en phrases complètes. Ne t’excuse pas inutilement et réponds directement à la question."""
+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, history=[]):
-    # Limiter l'historique pour accélérer
-    history = history[-3:]
-    
-    context = "\n".join([f"Utilisateur: {m[0]}\nAria: {m[1]}" for m in history])
-    prompt = f"{system_prompt}\n{context}\nUtilisateur: {message}\nAria:"
+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=60,  # plus court pour CPU
+        max_new_tokens=250,   # plus long pour éviter les coupures
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
         repetition_penalty=1.1
     )[0]["generated_text"]
-    
-    reply = resp.split("Aria:")[-1].strip()
-    history.append([message, reply])
-    
-    # Convertir en tuples pour Gradio
-    chat_display = [(m[0], m[1]) for m in history]
-    return chat_display, history
+
+    # 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:
-    chat_ui = gr.Chatbot()
+    chatbot_ui = gr.Chatbot()
+    state = gr.State([])  # sauvegarde de l'historique
     msg = gr.Textbox(placeholder="Écris un message...")
-    msg.submit(chat, [msg, chat_ui], [chat_ui, chat_ui])
+
+    msg.submit(chat, [msg, state], [chatbot_ui, state])
+    msg.submit(lambda: "", None, msg)  # reset input après envoi
 
 demo.launch()