Update app.py

app.py

@@ -1,34 +1,50 @@
 import os
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import torch
 
-# Récupérer le token Hugging Face depuis les secrets
+# Récupérer le token depuis les secrets/variables d'environnement
 hf_token = os.environ.get("HF_TOKEN")
 
-# Modèle léger et rapide
-model_name = "tiiuae/phi-4-mini"
+# Nom du modèle MLC quantifié (rapide et léger)
+model_name = "mlc-ai/gemma-2b-it-q4f16_1-MLC"
 
-# Charger le tokenizer et le modèle avec token
+print("Chargement du tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
+
+print("Chargement du modèle...")
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype=torch.float16,
-    device_map="auto",
+    torch_dtype=torch.float16,   # Float16 pour accélérer si GPU
+    device_map="auto",           # CPU ou GPU si disponible
     use_auth_token=hf_token
 )
 
 # Prompt système pour Aria
-system_prompt = """Tu es Aria, une IA bienveillante et polie qui répond de façon concise et claire."""
+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=[]):
+    # Construire le prompt avec l'historique
     prompt = system_prompt + "\n" + "\n".join([f"Utilisateur: {m[0]}\nAria: {m[1]}" for m in history]) + f"\nUtilisateur: {message}\nAria:"
+    
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=200)
+    
+    # Génération rapide
+    generation_config = GenerationConfig(
+        max_new_tokens=200,
+        temperature=0.7,
+        do_sample=True,
+        top_p=0.9
+    )
+    
+    outputs = model.generate(**inputs, generation_config=generation_config)
     reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
     # Récupérer uniquement la réponse d'Aria
     reply = reply.split("Aria:")[-1].strip()
     history.append((message, reply))
+    
     return reply, history
 
 # Interface Gradio