translate-package 0.0.8__tar.gz → 0.1.0__tar.gz

{translate_package-0.0.8 → translate_package-0.1.0}/PKG-INFO

@@ -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 → translate_package-0.1.0}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="translate_package",
-    version="0.0.8",
+    version="0.1.0",
     author="Oumar Kane",
     packages=find_packages(),
     author_email="oumar.kane@univ-thies.sn",

{translate_package-0.0.8 → translate_package-0.1.0}/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-0.0.8 → translate_package-0.1.0}/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-0.0.8 → translate_package-0.1.0}/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 → translate_package-0.1.0}/translate_package.egg-info/PKG-INFO

@@ -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