Update app.py

app.py

@@ -1,31 +1,39 @@
-import os
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import os
 
-# Token Hugging Face stocké dans les secrets
-hf_token = os.getenv("HF_TOKEN")
+# Récupérer le token Hugging Face (secret)
+hf_token = os.environ.get("HF_TOKEN")
 
-# Modèle réel et rapide
-model_name = "TheBloke/Guanaco-3B-Uncensored-v2-GPTQ"
+# Nom du modèle léger et rapide
+model_name = "tiiuae/gemma-2b"
 
-# Charger tokenizer et modèle avec le token
-tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
+# Charger le tokenizer et le modèle
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    token=hf_token
+    torch_dtype=torch.float16,  # moins lourd
+    device_map="auto",           # utilise GPU si dispo sinon CPU
+    use_auth_token=hf_token
 )
 
 # Prompt système pour Aria
-system_prompt = """Tu es Aria, une IA bienveillante et concise."""
+system_prompt = """Tu es Aria, une IA bienveillante et polie qui répond de façon concise et claire."""
 
 def chat(message, history=[]):
+    # Construire le prompt complet avec 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=150)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=200,
+        do_sample=True,          # rend les réponses plus naturelles
+        temperature=0.7
+    )
     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