Update app.py

app.py

@@ -1,38 +1,30 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Modèle léger pour CPU
-model_name = "OpenAssistant/replit-1B"
+MODEL = "prithivMLmods/Llama-SmolTalk-3.2-1B-Instruct"
 
-# Charger le tokenizer et le modèle
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+model = AutoModelForCausalLM.from_pretrained(MODEL, device_map="auto")
 
-# Créer le pipeline pour la génération
 chatbot = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=-1,  # CPU
+    device_map="auto",
 )
 
-# Prompt système pour Aria
 system_prompt = "Tu es Aria, une IA bienveillante et polie qui répond de façon concise et claire."
 
-# Fonction de chat
 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:"
-    
-    reply = chatbot(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]["generated_text"]
-    
-    # Ne garder que la réponse d'Aria (après le dernier "Aria:")
-    reply = reply.split("Aria:")[-1].strip()
-    
+
+    resp = chatbot(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]["generated_text"]
+    reply = resp.split("Aria:")[-1].strip()
+
     history.append((message, reply))
     return reply, history
 
-# Interface Gradio
 with gr.Blocks() as demo:
     chat_ui = gr.Chatbot()
     msg = gr.Textbox(placeholder="Écris un message...")