Update app.py

app.py

@@ -11,21 +11,11 @@ import json
 import os
 
 model_name = "Woziii/llama-3-8b-chat-me"
-
-# Initialiser le tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Configurer le pad token
-tokenizer.pad_token = tokenizer.eos_token
-
-# Initialiser le modèle
 model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
-
-# Configurer le pad token ID du modèle
-model.config.pad_token_id = tokenizer.pad_token_id
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 MAX_MAX_NEW_TOKENS = 250
-DEFAULT_MAX_NEW_TOKENS = 50
+DEFAULT_MAX_NEW_TOKENS = 70
 MAX_INPUT_TOKEN_LENGTH = 2048
 
 
@@ -159,53 +149,17 @@ def truncate_to_questions(text, max_questions):
     
     return ' '.join(truncated_sentences)
 
-
-
-
-def find_logical_stop(text, max_tokens):
-    """
-    Trouve un point d'arrêt logique dans le texte, sans dépasser max_tokens.
-    """
-    # Définir les motifs de fin logiques
-    end_patterns = r'(?<=[.!?])\s+|\n|\. |\! |\? '
-    
-    # Diviser le texte en segments logiques
-    segments = re.split(end_patterns, text)
-    
-    current_length = 0
-    result = ""
-    
-    for segment in segments:
-        segment_tokens = len(segment.split())
-        if current_length + segment_tokens <= max_tokens:
-            result += segment + " "
-            current_length += segment_tokens
-        else:
-            break
-    
-    # Nettoyer et finaliser le résultat
-    result = result.strip()
-    if result and result[-1] not in ".!?":
-        result += "."
-    
-    return result
-
 def post_process_response(response, is_short_response, max_questions=2):
-    # Limiter au nombre spécifié de questions
+    # Limiter au nombre spécifié de questions, quelle que soit la longueur de la réponse
     truncated_response = truncate_to_questions(response, max_questions)
     
-    # Déterminer la limite de tokens en fonction du type de réponse
+    # Appliquer la limitation de longueur si nécessaire pour les réponses courtes
     if is_short_response:
-        max_tokens = 70
-    else:
-        max_tokens = 150  # Ajustez selon vos besoins
+        sentences = re.split(r'(?<=[.!?])\s+', truncated_response)
+        if len(sentences) > 2:
+            return ' '.join(sentences[:2]).strip()
     
-    # Trouver un point d'arrêt logique
-    final_response = find_logical_stop(truncated_response, max_tokens)
-    
-    return final_response
-
-
+    return truncated_response.strip()
     
 def check_coherence(response):
     sentences = re.split(r'(?<=[.!?])\s+', response)
@@ -214,7 +168,7 @@ def check_coherence(response):
         return False
     return True
 
-@spaces.GPU(duration=180)
+@spaces.GPU(duration=120)
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
@@ -237,73 +191,57 @@ def generate(
     if response_type == "short":
         max_new_tokens = max(70, max_new_tokens)
     elif response_type == "long":
-        max_new_tokens = min(max(200, max_new_tokens), 250)
+        max_new_tokens = min(max(200, max_new_tokens), 300)
     else:  # medium
-        max_new_tokens = min(max(70, max_new_tokens), 150)
+        max_new_tokens = min(max(100, max_new_tokens), 150)
 
-    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[-5:]:
-            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")
-        attention_mask = input_ids.ne(tokenizer.pad_token_id).long()
+    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})
 
-        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.")
+    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 = 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)
+        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
         
-        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 partial_output
 
-        t = Thread(target=model.generate, kwargs=generate_kwargs)
-        t.start()
-
-        outputs = []
-        for text in streamer:
-            outputs.append(text)
-            current_output = "".join(outputs)
-            processed_output = post_process_response(current_output, response_type == "short")
-            
-            if response_type == "long" and not check_coherence(processed_output):
-                yield "Je m'excuse, ma réponse manquait de cohérence. Pouvez-vous reformuler votre question ?"
-                return
-            
-            yield processed_output
-
-        final_output = post_process_response("".join(outputs), response_type == "short")
-        yield final_output
-
-    except Exception as e:
-        print(f"Une erreur s'est produite : {str(e)}")
-        yield "Désolé, une erreur s'est produite. Veuillez réessayer."
+    yield post_process_response("".join(outputs), response_type == "short")
 
-    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 ""
@@ -449,4 +387,4 @@ N'hésitez pas à aborder des sujets variés, allant de l'intelligence artificie
     chat_interface.render()
     chat_interface.chatbot.like(vote, [chat_interface.chatbot], None)
 
-demo.queue(max_size=20, default_concurrency_limit=2).launch(max_threads=10)
+demo.queue(max_size=20, default_concurrency_limit=2).launch(max_threads=10)