translate-package 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

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.dist-info → translate_package-0.0.4.dist-info}/METADATA

@@ -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.dist-info → translate_package-0.0.4.dist-info}/RECORD

@@ -5,13 +5,13 @@ translate_package/errors/__init__.py,sha256=gu6XjAIghG4lLkYo8x_7_yyLRtK2FIvmC-Wc
 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=5QQpjs_HR9mnPryMyfYpcMgU5tHAAj-eVrv3oGmjR5Y,9963
+translate_package/models/machine_translation.py,sha256=LW1qNVAP-c2rKYC58L6cPQ4LU2ePD8wg_Z8CcHSI8MY,10575
 translate_package/tokenization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 translate_package/tokenization/load_tokenizer.py,sha256=vzCHS0ZDSJyr0y08zNvupMtD2jP8A16EBN-ob0LJHG0,1344
 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.2.dist-info/METADATA,sha256=OikhBGLwAfDae4ZhSbzNNdHVfj08UxIjSR_P6KryzMA,860
-translate_package-0.0.2.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-translate_package-0.0.2.dist-info/top_level.txt,sha256=8e2HIrGAMzoSukqu2q929dOJMV1zGYKI_BAFwl-P7XU,18
-translate_package-0.0.2.dist-info/RECORD,,
+translate_package-0.0.4.dist-info/METADATA,sha256=rT99Y10OFAqYCHBsjQiXxydlQhcFS2Ak-zLheAQJEE0,860
+translate_package-0.0.4.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+translate_package-0.0.4.dist-info/top_level.txt,sha256=8e2HIrGAMzoSukqu2q929dOJMV1zGYKI_BAFwl-P7XU,18
+translate_package-0.0.4.dist-info/RECORD,,