Update app.py

app.py

@@ -3,29 +3,47 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 MODEL = "prithivMLmods/Llama-SmolTalk-3.2-1B-Instruct"
 
+# Charger le tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model = AutoModelForCausalLM.from_pretrained(MODEL, device_map="auto")
 
+# Charger le modèle en 8 bits pour accélérer
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL,
+    device_map="auto",
+    load_in_8bit=True  # optimisation vitesse / mémoire
+)
+
+# Pipeline avec paramètres par défaut optimisés
 chatbot = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device_map="auto",
+    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."
 
+# Fonction de chat optimisée
 def chat(message, history=[]):
     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=150, do_sample=True, temperature=0.7)[0]["generated_text"]
-    reply = resp.split("Aria:")[-1].strip()
+    # Paramètres réduits pour accélérer la génération
+    resp = chatbot(
+        prompt,
+        max_new_tokens=60,       # Limite pour réduire temps de calcul
+        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, history
 
-
+# Interface Gradio
 with gr.Blocks() as demo:
     chat_ui = gr.Chatbot()
     msg = gr.Textbox(placeholder="Écris un message...")