Update app.py

app.py

@@ -3,16 +3,22 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 
-# 🔹 Token Hugging Face
+# 🔹 Récupérez le token depuis les secrets du Space
 hf_token = os.environ.get("HF_TOKEN")
+
 if hf_token is None:
-    raise ValueError("⚠️ HF_TOKEN manquant dans les secrets du Space.")
+    raise ValueError("⚠️ Le token Hugging Face (HF_TOKEN) n'est pas trouvé. "
+                     "Vérifie que tu l’as bien ajouté dans les secrets du Space.")
 
-# 🔹 ID du modèle
+# 🔹 Charger le tokenizer et le modèle
 model_id = "SafaaAI/final_llm_darija_fr_tech"
 
-# 🔹 Charger tokenizer + modèle (forcé en text-only)
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(
+    model_id,
+    token=hf_token,
+    trust_remote_code=True
+)
+
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     token=hf_token,
@@ -20,35 +26,60 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-# 🔹 Fonction de génération texte
+# 🔹 Fonction d'inférence
 def chat_with_model(message, history):
     history = history or []
-    inputs = tokenizer(message, return_tensors="pt").to(model.device)
-
+    full_prompt = "A chat between a curious user and an AI assistant."
+    
+    # Construire le prompt manuellement
+    for user_message, bot_message in history:
+        full_prompt += f" USER: {user_message} ASSISTANT: {bot_message}"
+    
+    # Ajouter le message actuel de l'utilisateur
+    full_prompt += f" USER: {message} ASSISTANT:"
+    
+    # Encoder le prompt avec le tokenizer
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+    
+    # 📝 Extraire explicitement input_ids et attention_mask
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+    
+    # Générer la réponse
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
+        output_ids = model.generate(
+            input_ids,
+            attention_mask=attention_mask,
             max_new_tokens=200,
             do_sample=True,
-            temperature=0.7,
-            top_p=0.9
+            top_p=0.9,
+            temperature=0.7
         )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Décoder la sortie
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    
+    # Nettoyer la réponse pour ne pas inclure le prompt
+    response_start_index = response.rfind("ASSISTANT:")
+    if response_start_index != -1:
+        response = response[response_start_index + len("ASSISTANT:"):].strip()
+    
     history.append((message, response))
     return history, history
 
-# 🔹 UI Gradio
+# 🔹 Interface Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## 💬 SafaaAI - Multimodal LLM (mode texte seulement)")
-    chatbot = gr.Chatbot(type="messages")
+    gr.Markdown("## 💬 Chatbot SafaaAI - LLM (Darija + Français + Technique)")
+    
+    chatbot = gr.Chatbot()
     msg = gr.Textbox(label="Écris ton message ici")
-    clear = gr.Button("🧹 Effacer")
+    clear = gr.Button("🧹 Effacer la conversation")
 
     state = gr.State([])
 
     msg.submit(chat_with_model, [msg, state], [chatbot, state])
     clear.click(lambda: ([], []), None, [chatbot, state])
 
+# 🔹 Lancer l'application
 if __name__ == "__main__":
     demo.launch()