translate-package 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

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

@@ -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
@@ -28,4 +28,5 @@ Requires-Dist: selenium
 Requires-Dist: sentencepiece
 Requires-Dist: peft
 Requires-Dist: rouge-score
+Requires-Dist: wolof-translate (==0.0.1)
 

{translate_package-0.0.1.dist-info → translate_package-0.0.3.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=fk5WrxX_8MsQiVy318yQdTBiH4oK1dLmqU6sAyY0_4Y,10438
 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.1.dist-info/METADATA,sha256=eHD4CGRhoQHOl-EXUPYPTQ9kGQCw5q3Ju0xkmB3PGTM,819
-translate_package-0.0.1.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-translate_package-0.0.1.dist-info/top_level.txt,sha256=8e2HIrGAMzoSukqu2q929dOJMV1zGYKI_BAFwl-P7XU,18
-translate_package-0.0.1.dist-info/RECORD,,
+translate_package-0.0.3.dist-info/METADATA,sha256=UyqJD8nfZIttZRMRLDJMYN68WZy9jdWwkt8pQiSXzaE,860
+translate_package-0.0.3.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+translate_package-0.0.3.dist-info/top_level.txt,sha256=8e2HIrGAMzoSukqu2q929dOJMV1zGYKI_BAFwl-P7XU,18
+translate_package-0.0.3.dist-info/RECORD,,