Update goai_helpers/goai_traduction.py

goai_helpers/goai_traduction.py

@@ -2,7 +2,7 @@ import torch
 import spaces
 import re
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from goai_helpers.utils import MooreConverter
+from goai_helpers.utils import MooreConverter, mark_numbers, unmark_numbers
 from huggingface_hub import login
 import os
 
@@ -62,6 +62,8 @@ def translate_chunk(text, src_lang, tgt_lang):
     else:
         model_id = "ArissBandoss/nllb-200-3.3B-fr2mos"
         #model_id = "ArissBandoss/nllb-200-3.3B-mos-fr-bidirectional-peft"
+
+    text = mark_numbers(text)
     
     tokenizer = AutoTokenizer.from_pretrained(model_id, token=auth_token)
     model = AutoModelForSeq2SeqLM.from_pretrained(model_id, token=auth_token).to(device)
@@ -86,6 +88,15 @@ def translate_chunk(text, src_lang, tgt_lang):
     # Décodage
     translation = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
 
+    translation = unmark_numbers(translation)
+
+    number_converter = MooreConverter()
+    numbers = re.findall(r'\b\d+\b', translation)
+    for number in numbers:
+        moore_number = number_converter.number_to_moore(int(number))
+        if moore_number:  # Only replace if conversion succeeded
+            translation = translation.replace(number, moore_number)
+
     
     return translation