translate-package 0.0.1__tar.gz → 0.0.3__tar.gz

{translate_package-0.0.1 → translate_package-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: translate_package
-Version: 0.0.1
+Version: 0.0.3
 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.1 → translate_package-0.0.3}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="translate_package",
-    version="0.0.1",
+    version="0.0.3",
     author="Oumar Kane",
     packages=find_packages(),
     author_email="oumar.kane@univ-thies.sn",
@@ -32,6 +32,7 @@ setup(
         "sentencepiece",
         "peft",
         "rouge-score",
-        "sacrebleu"
+        "sacrebleu",
+        "wolof-translate==0.0.1"
     ],
 )

{translate_package-0.0.1 → translate_package-0.0.3}/translate_package/models/machine_translation.py

@@ -263,17 +263,32 @@ class MachineTranslationTransformer(pl.LightningModule):
         )
 
         # generate predictions
-        predictions = self.model.generate(
-            input_ids=batch["input_ids"],
-            attention_mask=batch["attention_mask"],
-            max_new_tokens=self.max_new_tokens,
-            do_sample=self.num_beams > 0,
-            num_beams=self.num_beams
-        ) if not self.model_generation in ["lstm"] else self.model.generate(
-            input=batch["input_ids"],
-            max_new_tokens=self.max_new_tokens,
-            use_sampling=True
-        )
+        if not self.model_generation in ["lstm"] and self.num_beams > 0:
+            
+            predictions = self.model.generate(
+                input_ids=batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                max_new_tokens=self.max_new_tokens,
+                do_sample=True,
+                num_beams=self.num_beams
+            )  
+            
+        elif not self.model_generation in ["lstm"]:
+            
+            predictions = self.model.generate(
+                input_ids=batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                max_new_tokens=self.max_new_tokens,
+                do_sample=False
+            ) 
+        
+        else:
+            
+            predictions = self.model.generate(
+                input=batch["input_ids"],
+                max_new_tokens=self.max_new_tokens,
+                use_sampling=True
+            )
 
         # decode the labels
         predictions = self.tokenizer.batch_decode(predictions, skip_special_tokens=True)

{translate_package-0.0.1 → translate_package-0.0.3}/translate_package.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: translate-package
-Version: 0.0.1
+Version: 0.0.3
 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.1 → translate_package-0.0.3}/translate_package.egg-info/requires.txt

@@ -23,3 +23,4 @@ sentencepiece
 peft
 rouge-score
 sacrebleu
+wolof-translate==0.0.1