Update app.py

app.py

@@ -1,51 +1,69 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from threading import Lock
 
-# Charger le modèle et le tokenizer
+# Chargement du modèle
 model_name = "Dorian2B/Vera-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
+model.eval()
+
+# Verrou pour éviter les conflits de threads
+generate_lock = Lock()
 
-def generate_response(message, history):
-    # Formatage de l'historique de conversation
+def format_prompt(history, new_message):
+    """Formate l'historique et le nouveau message pour le modèle."""
     prompt = ""
-    if history:
-        for user_msg, bot_msg in history:
-            prompt += f"<|user|>{user_msg}</s>\n<|assistant|>{bot_msg}</s>\n"
-    prompt += f"<|user|>{message}</s>\n<|assistant|>"
-    
-    # Encodage et génération
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
-        repetition_penalty=1.1,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    
-    # Décodage et nettoyage de la réponse
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = response.split("<|assistant|>")[-1].strip()
-    
-    return response
-
-# Interface Gradio avec ChatInterface
+    for user_msg, bot_msg in history:
+        prompt += f"<|user|>{user_msg}</s>\n<|assistant|>{bot_msg}</s>\n"
+    prompt += f"<|user|>{new_message}</s>\n<|assistant|>"
+    return prompt
+
+def generate_stream(history, new_message):
+    """Génère une réponse en streaming avec contexte."""
+    prompt = format_prompt(history, new_message)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    # Génération en streaming
+    with generate_lock:
+        with torch.no_grad():
+            for chunk in model.generate(
+                **inputs,
+                max_new_tokens=1024,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                repetition_penalty=1.1,
+                eos_token_id=tokenizer.eos_token_id,
+                streamer=None,  # (Remplacez par un vrai streamer si disponible)
+            ):
+                decoded = tokenizer.decode(chunk[0], skip_special_tokens=True)
+                if decoded.startswith(prompt):  # Supprime le prompt
+                    decoded = decoded[len(prompt):]
+                yield decoded.strip()
+
+def chat_interface(message, history):
+    """Fonction pour Gradio ChatInterface."""
+    full_response = ""
+    for chunk in generate_stream(history, message):
+        full_response += chunk
+        yield full_response
+
+# Interface Gradio
 demo = gr.ChatInterface(
-    fn=generate_response,
-    title="Chat avec Vera-Instruct",
-    description="Discutez avec le modèle Vera-Instruct de Dorian2B",
-    examples=["Bonjour, comment ça va ?", 
-              "Explique-moi le théorème de Pythagore",
-              "Donne-moi une recette de cookies"],
-    theme="soft"
+    fn=chat_interface,
+    title="💬 Vera-Instruct Chat (avec Contexte & Streaming)",
+    description="Discutez avec le modèle **Dorian2B/Vera-Instruct**.<br>Le modèle conserve le contexte de la conversation.",
+    examples=["Bonjour ! Comment vas-tu ?", "Explique-moi l'IA générative."],
+    theme="soft",
+    retry_btn=None,
+    undo_btn=None,
 )
 
-# Pour Hugging Face Spaces, utilisez launch()
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch(debug=True)