Update goai_helpers/goai_traduction.py

goai_helpers/goai_traduction.py

@@ -9,14 +9,27 @@ max_length = 512
 auth_token = os.getenv('HF_SPACE_TOKEN')
 login(token=auth_token)
 
-def split_text_intelligently(text, max_chunk_length=100):
+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):
     """
-    Divise le texte en chunks en respectant les phrases complètes.
+    Divise le texte en chunks en respectant les phrases et en comptant les tokens.
     """
-    # Séparation basée sur les phrases (utilise les points, points d'interrogation, etc.)
+    tokenizer = get_tokenizer(src_lang, tgt_lang)
+    tokenizer.src_lang = src_lang
+    
+    # Séparation basée sur les phrases
     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
@@ -25,12 +38,17 @@ def split_text_intelligently(text, max_chunk_length=100):
         else:
             sentence = sentences[i]
         
-        # 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:
+        # 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:
             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:
@@ -39,10 +57,15 @@ def split_text_intelligently(text, max_chunk_length=100):
     return chunks
 
 @spaces.GPU
-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)
+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)
         translations = []
         
         for chunk in chunks:
@@ -73,14 +96,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 pour éviter les répétitions
+    # Paramètres de génération optimisés
     outputs = model.generate(
         **inputs,
         forced_bos_token_id=tgt_lang_id,
         max_new_tokens=512,
         num_beams=5,
-        no_repeat_ngram_size=4,
-        repetition_penalty=2.0,
+        no_repeat_ngram_size=3,
+        repetition_penalty=1.5,
         length_penalty=1.0,
         early_stopping=True
     )
@@ -91,4 +114,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)
+    return goai_traduction(input_text, src_lang, tgt_lang, max_tokens_per_chunk=200)