translate-package 0.0.8__py3-none-any.whl → 0.1.0__py3-none-any.whl

translate_package/__init__.py

@@ -14,7 +14,7 @@ from transformers import (
     Trainer, AutoModelForSeq2SeqLM,
     get_linear_schedule_with_warmup,
     T5ForConditionalGeneration, Adafactor, BartForConditionalGeneration,
-    MT5ForConditionalGeneration, AdamWeightDecay
+    MT5ForConditionalGeneration, AdamWeightDecay, AutoTokenizer
 )
 from wolof_translate.utils.bucket_iterator import SequenceLengthBatchSampler, BucketSampler
 from wolof_translate.utils.sent_transformers import TransformerSequences

translate_package/models/machine_translation.py

@@ -11,6 +11,7 @@ from translate_package import (
     MT5ForConditionalGeneration,
     BartForConditionalGeneration,
     AutoModelForSeq2SeqLM,
+    AutoTokenizer,
     Adafactor,
     AdamWeightDecay
 )
@@ -57,7 +58,7 @@ class MachineTranslationTransformer(pl.LightningModule):
         hidden_size=128,
         dropout=0.1,
         bidirectional=False,
-        length_penalty=1.2
+        length_penalty=1.2,
     ):
 
         super().__init__()
@@ -157,7 +158,7 @@ class MachineTranslationTransformer(pl.LightningModule):
                 warmup_init = False
             )
         
-        elif self.model_generation in ["bart", "mbart"]:
+        elif self.model_generation in ["bart", "mbart", "nllb"]:
             
             optimizer = torch.optim.AdamW(
                 self.parameters(), lr=self.lr, weight_decay=self.weight_decay
@@ -169,11 +170,11 @@ class MachineTranslationTransformer(pl.LightningModule):
                 self.parameters(), lr=self.lr, weight_decay=self.weight_decay
             )
 
-        if self.model_generation in ["t5", "lstm", "mt5"]:
+        if self.model_generation in ["t5", "lstm"]:
             
             return [optimizer]
 
-        elif self.model_generation in ["bart"]:
+        elif self.model_generation in ["bart", "mt5"]:
 
             scheduler = get_linear_schedule_with_warmup(
                 optimizer,

translate_package/tokenization/load_tokenizer.py

@@ -3,13 +3,18 @@ from transformers import T5TokenizerFast
 from transformers import AutoTokenizer
 import os
 
-def load_tokenizer(tokenizer_name, model, dir_path, file_name, model_name = None):
+BCP_47_languages = {
+    'french': 'fra_Latn',
+    'wolof': 'wol_Latn',
+}
+
+def load_tokenizer(tokenizer_name, model, dir_path, file_name, model_name = None, src_lang = "french", tgt_lang = "wolof"):
     
     if model == "nllb":
         
         if not model_name is None:
         
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang = BCP_47_languages[src_lang], tgt_lang = BCP_47_languages[tgt_lang])
             
             print(f"The {model}'s tokenizer was successfully loaded")
         

{translate_package-0.0.8.dist-info → translate_package-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: translate-package
-Version: 0.0.8
+Version: 0.1.0
 Summary: Contain functions and classes to efficiently train a sequence to sequence to translate between two languages.
 Author: Oumar Kane
 Author-email: oumar.kane@univ-thies.sn

{translate_package-0.0.8.dist-info → translate_package-0.1.0.dist-info}/RECORD

@@ -1,17 +1,17 @@
-translate_package/__init__.py,sha256=Nckjm15LBEfKSU5-EBjfRkHJQhWELzwow0BD3rKtmkw,1297
+translate_package/__init__.py,sha256=miie3aAeUYHsVk2O-kd4T86fFksuCiY70Eo6RNeY1Oo,1312
 translate_package/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 translate_package/data/data_preparation.py,sha256=AQFmIJawPL5kaWGhTGPCfizIQD00XqQ-cANaICZ-1Ow,14882
 translate_package/errors/__init__.py,sha256=gu6XjAIghG4lLkYo8x_7_yyLRtK2FIvmC-WcfJaeOlg,299
 translate_package/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 translate_package/models/gradient_observation.py,sha256=P91UA5i-RdkK46TqpPOJ54DsUYgTI9cRohgPS1Ch0Lc,294
 translate_package/models/lstm.py,sha256=OPkvvceowz5JqdGGH4cfPhH23kbP11z-29zIJn5d8ig,3273
-translate_package/models/machine_translation.py,sha256=j8ZXh9UEElixwy7RwBFE_sO0EZrkaMOGs7W39EWvXhc,10575
+translate_package/models/machine_translation.py,sha256=a9YOEgsuGnkAvaEv8jML95YESNgkHu96ZXHm6_esv44,10604
 translate_package/tokenization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-translate_package/tokenization/load_tokenizer.py,sha256=vzCHS0ZDSJyr0y08zNvupMtD2jP8A16EBN-ob0LJHG0,1344
+translate_package/tokenization/load_tokenizer.py,sha256=g8j5pDmimFhwjpeYNkWot0hXMzAqqURbtedcQK-1xYE,1543
 translate_package/tokenization/train_tokenizer.py,sha256=RkdT5DUx201OBNaswM6m54iqcrmCThd3ITLguQb_zVM,3347
 translate_package/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 translate_package/utils/checkpoint.py,sha256=GqymRvF8_QZgrQq9m79Ppj6Qr7NQm78kDARm3p_chC0,322
-translate_package-0.0.8.dist-info/METADATA,sha256=GiNYpQ6GrFumBn--bFa0cFijKTYjoMV2Wib8cm7R6Cg,850
-translate_package-0.0.8.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-translate_package-0.0.8.dist-info/top_level.txt,sha256=8e2HIrGAMzoSukqu2q929dOJMV1zGYKI_BAFwl-P7XU,18
-translate_package-0.0.8.dist-info/RECORD,,
+translate_package-0.1.0.dist-info/METADATA,sha256=l90EsBgKdqObcW_Q8w-2A3UpM-ji8Yn8V3QyLxUg-gk,850
+translate_package-0.1.0.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+translate_package-0.1.0.dist-info/top_level.txt,sha256=8e2HIrGAMzoSukqu2q929dOJMV1zGYKI_BAFwl-P7XU,18
+translate_package-0.1.0.dist-info/RECORD,,