translate-package 0.0.2__tar.gz → 0.0.4__tar.gz

{translate_package-0.0.2 → translate_package-0.0.4}/PKG-INFO

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

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

{translate_package-0.0.2 → translate_package-0.0.4}/translate_package/models/machine_translation.py

@@ -48,7 +48,7 @@ class MachineTranslationTransformer(pl.LightningModule):
         lora_alpha=32,
         lora_dropout=0.05,
         bias="none",
-        max_new_tokens=200,
+        max_new_tokens=90,
         predict_with_generate=True,
         num_beams=0,
         use_peft=False,
@@ -56,7 +56,8 @@ class MachineTranslationTransformer(pl.LightningModule):
         num_layers=6,
         hidden_size=128,
         dropout=0.1,
-        bidirectional=False
+        bidirectional=False,
+        length_penalty=1.2
     ):
 
         super().__init__()
@@ -131,6 +132,8 @@ class MachineTranslationTransformer(pl.LightningModule):
 
         self.num_beams = num_beams
         
+        self.length_penalty = length_penalty
+        
         self.model_generation = model_generation
 
         self.predictions = {
@@ -263,17 +266,33 @@ 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,
+                length_penalty=self.length_penalty
+            )  
+            
+        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.2 → translate_package-0.0.4}/translate_package.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: translate-package
-Version: 0.0.2
+Version: 0.0.4
 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