Update app.py

app.py

@@ -1,33 +1,55 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-MODEL = "gpt2"
+MODEL = "prithivMLmods/Llama-SmolTalk-3.2-1B-Instruct"
 
+# Charger le tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model = AutoModelForCausalLM.from_pretrained(MODEL)
 
+# Charger le modèle en 8 bits pour accélérer et réduire l’usage mémoire
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL,
+    device_map="auto",
+    load_in_8bit=True  # accélère sur CPU/peu de mémoire
+)
+
+# Pipeline de génération
 chatbot = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=-1  # CPU
+    device_map="auto"
 )
 
+# Prompt système
 system_prompt = "Tu es Aria, une IA bienveillante et polie qui répond de façon concise et claire."
 
 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:"
 
-    resp = chatbot(prompt, max_new_tokens=50, do_sample=True, temperature=0.7)[0]["generated_text"]
-    reply = resp.split("Aria:")[-1].strip()
+    resp = chatbot(
+        prompt,
+        max_new_tokens=60,  # plus court pour CPU
+        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])
-    return history
+    
+    # Convertir en tuples pour Gradio
+    chat_display = [(m[0], m[1]) for m in history]
+    return chat_display, history
 
 with gr.Blocks() as demo:
     chat_ui = gr.Chatbot()
     msg = gr.Textbox(placeholder="Écris un message...")
-    msg.submit(chat, [msg, chat_ui], [chat_ui])
+    msg.submit(chat, [msg, chat_ui], [chat_ui, chat_ui])
 
 demo.launch()