Update goai_helpers/goai_traduction.py

goai_helpers/goai_traduction.py

@@ -9,27 +9,14 @@ max_length = 512
 auth_token = os.getenv('HF_SPACE_TOKEN')
 login(token=auth_token)
 
-def get_tokenizer(src_lang, tgt_lang):
-    """Initialise et retourne le tokenizer approprié"""
-    if src_lang == "mos_Latn" and tgt_lang == "fra_Latn":
-        model_id = "ArissBandoss/3b-new-400"
-    else:
-        model_id = "ArissBandoss/nllb-200-distilled-600M-finetuned-fr-to-mos-V4"
-    
-    return AutoTokenizer.from_pretrained(model_id, token=auth_token)
-
-def split_text_by_tokens(text, src_lang, tgt_lang, max_tokens_per_chunk=200):
+def split_text_intelligently(text, max_chunk_length=100):
     """
-    Divise le texte en chunks en respectant les phrases et en comptant les tokens.
+    Divise le texte en chunks en respectant les phrases complètes.
     """
-    tokenizer = get_tokenizer(src_lang, tgt_lang)
-    tokenizer.src_lang = src_lang
-    
-    # Séparation basée sur les phrases
+    # Séparation basée sur les phrases (utilise les points, points d'interrogation, etc.)
     sentences = re.split(r'([.!?])', text)
     chunks = []
     current_chunk = ""
-    current_tokens = 0
     
     for i in range(0, len(sentences), 2):
         # Reconstruire la phrase avec sa ponctuation
@@ -38,17 +25,12 @@ def split_text_by_tokens(text, src_lang, tgt_lang, max_tokens_per_chunk=200):
         else:
             sentence = sentences[i]
         
-        # Calculer le nombre de tokens pour cette phrase
-        sentence_tokens = len(tokenizer.encode(sentence))
-        
-        # Si l'ajout de cette phrase dépasse la limite de tokens, on crée un nouveau chunk
-        if current_tokens + sentence_tokens > max_tokens_per_chunk and current_chunk:
+        # Si l'ajout de cette phrase dépasse la longueur maximale, on crée un nouveau chunk
+        if len(current_chunk) + len(sentence) > max_chunk_length and current_chunk:
             chunks.append(current_chunk.strip())
             current_chunk = sentence
-            current_tokens = sentence_tokens
         else:
             current_chunk += sentence
-            current_tokens += sentence_tokens
     
     # Ajouter le dernier chunk s'il reste du texte
     if current_chunk:
@@ -57,15 +39,10 @@ def split_text_by_tokens(text, src_lang, tgt_lang, max_tokens_per_chunk=200):
     return chunks
 
 @spaces.GPU
-def goai_traduction(text, src_lang, tgt_lang, max_tokens_per_chunk=200):
-    # Vérifier si le texte a besoin d'être divisé en comptant les tokens
-    tokenizer = get_tokenizer(src_lang, tgt_lang)
-    tokenizer.src_lang = src_lang
-    
-    text_tokens = len(tokenizer.encode(text))
-    
-    if text_tokens > max_tokens_per_chunk:
-        chunks = split_text_by_tokens(text, src_lang, tgt_lang, max_tokens_per_chunk)
+def goai_traduction(text, src_lang, tgt_lang, max_chunk_length=100):
+    # Si le texte est trop long, le diviser en chunks
+    if len(text) > max_chunk_length:
+        chunks = split_text_intelligently(text, max_chunk_length)
         translations = []
         
         for chunk in chunks:
@@ -96,14 +73,14 @@ def translate_chunk(text, src_lang, tgt_lang):
     # ID du token de langue cible
     tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang)
     
-    # Paramètres de génération optimisés
+    # Paramètres de génération optimisés pour éviter les répétitions
     outputs = model.generate(
         **inputs,
         forced_bos_token_id=tgt_lang_id,
         max_new_tokens=512,
         num_beams=5,
-        no_repeat_ngram_size=3,
-        repetition_penalty=1.5,
+        no_repeat_ngram_size=4,
+        repetition_penalty=2.0,
         length_penalty=1.0,
         early_stopping=True
     )
@@ -114,4 +91,4 @@ def translate_chunk(text, src_lang, tgt_lang):
     return translation
 
 def real_time_traduction(input_text, src_lang, tgt_lang):
-    return goai_traduction(input_text, src_lang, tgt_lang, max_tokens_per_chunk=200)
+    return goai_traduction(input_text, src_lang, tgt_lang)