Update app.py

app.py

@@ -10,6 +10,9 @@ from huggingface_hub import HfApi, hf_hub_download
 import json
 import os
 
+tokenizer.pad_token = tokenizer.eos_token
+model.config.pad_token_id = tokenizer.pad_token_id
+
 model_name = "Woziii/llama-3-8b-chat-me"
 model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -168,7 +171,7 @@ def check_coherence(response):
         return False
     return True
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=180)
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -195,53 +198,67 @@ def generate(
     else:  # medium
         max_new_tokens = min(max(100, max_new_tokens), 150)
 
-    conversation = []
-    
-    # Ajout du system prompt et du LUCAS_KNOWLEDGE_BASE
-    enhanced_system_prompt = f"{system_prompt}\n\n{LUCAS_KNOWLEDGE_BASE}"
-    conversation.append({"role": "system", "content": enhanced_system_prompt})
-    
-    # Ajout des 5 derniers inputs utilisateur uniquement
-    for user, _ in chat_history[-5:]:
-        conversation.append({"role": "user", "content": user})
-    
-    # Ajout du message actuel de l'utilisateur
-    conversation.append({"role": "user", "content": message})
+    try:
+        conversation = []
+        
+        # Ajout du system prompt et du LUCAS_KNOWLEDGE_BASE
+        enhanced_system_prompt = f"{system_prompt}\n\n{LUCAS_KNOWLEDGE_BASE}"
+        conversation.append({"role": "system", "content": enhanced_system_prompt})
+        
+        # Ajout des 5 dernières interactions complètes (user uniquement)
+        for user, assistant in chat_history[-MAX_HISTORY_LENGTH:]:
+            conversation.append({"role": "user", "content": user})
+        
+        # Ajout du message actuel de l'utilisateur
+        conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
-    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"L'entrée de la conversation a été tronquée car elle dépassait {MAX_INPUT_TOKEN_LENGTH} tokens.")
-    
-    input_ids = input_ids.to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        temperature=temperature,
-        num_beams=1,
-    )
-
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        partial_output = post_process_response("".join(outputs), response_type == "short")
+        input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
+        attention_mask = input_ids.ne(tokenizer.pad_token_id).long()
+
+        if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+            input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+            attention_mask = attention_mask[:, -MAX_INPUT_TOKEN_LENGTH:]
+            gr.Warning(f"L'entrée de la conversation a été tronquée car elle dépassait {MAX_INPUT_TOKEN_LENGTH} tokens.")
         
-        if response_type == "long" and not check_coherence(partial_output):
-            yield "Je m'excuse, ma réponse manquait de cohérence. Pouvez-vous reformuler votre question ?"
-            return
+        input_ids = input_ids.to(model.device)
+        attention_mask = attention_mask.to(model.device)
+
+        streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
         
-        yield partial_output
+        generate_kwargs = dict(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            streamer=streamer,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            top_p=top_p,
+            temperature=temperature,
+            num_beams=1,
+        )
 
-    yield post_process_response("".join(outputs), response_type == "short")
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
+
+        outputs = []
+        for text in streamer:
+            outputs.append(text)
+            partial_output = post_process_response("".join(outputs), response_type == "short")
+            
+            if response_type == "long" and not check_coherence(partial_output):
+                yield "Je m'excuse, ma réponse manquait de cohérence. Pouvez-vous reformuler votre question ?"
+                return
+            
+            yield partial_output
+
+        yield post_process_response("".join(outputs), response_type == "short")
+
+    except Exception as e:
+        print(f"Une erreur s'est produite : {str(e)}")
+        yield "Désolé, une erreur s'est produite. Veuillez réessayer."
 
+    finally:
+        # Nettoyage de la mémoire GPU
+        torch.cuda.empty_cache()
 
 def vote(data: gr.LikeData, history):
     user_input = history[-1][0] if history else ""