Update finetune.py

finetune.py

@@ -26,7 +26,7 @@ def summarize_text_mt5(texts, model, tokenizer):
                        max_length=512, truncation=True, 
                        padding=True).to(model.device)
     summary_ids = model.generate(inputs.input_ids, 
-                                 max_length=128, 
+                                 max_length=60, 
                                  num_beams=4, length_penalty=2.0, 
                                  early_stopping=True)
     summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)
@@ -36,7 +36,7 @@ def summarize_text_mbart50(texts, model, tokenizer):
     inputs = tokenizer(texts, return_tensors="pt",
                        max_length=1024, truncation=True, 
                        padding=True).to(model.device)
-    summary_ids = model.generate(inputs.input_ids, max_length=128,
+    summary_ids = model.generate(inputs.input_ids, max_length=60,
                                  num_beams=4, length_penalty=2.0,
                                  early_stopping=True)
     summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)
@@ -94,10 +94,10 @@ def fine_tune(model_name, finetune_type, model, tokenizer, summarize_text, train
     print("Starting Fine-tuning...")
     if model_name == "mT5":
         max_input = 512 
-        max_output = 128
+        max_output = 60
     else:
         max_input = 1024
-        max_output = 128
+        max_output = 60
 
     train_dataset = train
     eval_dataset = val