Update finetune.py

finetune.py

@@ -13,7 +13,7 @@ from transformers import (AutoTokenizer, BitsAndBytesConfig, MBart50TokenizerFas
                           MBartForConditionalGeneration, TrainingArguments,
                           DataCollatorForSeq2Seq, EarlyStoppingCallback)
 from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
-# Get the absolute path of the current script
+
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
 MODELS = {
@@ -53,7 +53,6 @@ def experiments(model_name, finetune_type):
     """Runs an experiment with the given model and dataset."""
     print(f"Starting Experiment: on {model_name}")
 
-    # Construct dataset paths dynamically
     train = pd.read_csv(os.path.join(BASE_DIR, "datasets/train.csv"))
     train_fr = pd.read_csv(os.path.join(BASE_DIR, "datasets/train_fr.csv"))
     train_cross = pd.read_csv(os.path.join(BASE_DIR, "datasets/train_cross.csv"))
@@ -64,16 +63,6 @@ def experiments(model_name, finetune_type):
     test_fr = pd.read_csv(os.path.join(BASE_DIR, "datasets/test_fr.csv"))
     test_cross = pd.read_csv(os.path.join(BASE_DIR, "datasets/test_cross.csv"))
 
-    # print(len(train))
-    # print(len(train_fr))
-    # print(len(train_cross))
-    # print(len(val))
-    # print(len(val_fr))
-    # print(len(val_cross))
-    # print(len(test))
-    # print(len(test_fr))
-    # print(len(test_cross))
-
     model, tokenizer = download_model(model_name)
     print(f"Model {model_name} loaded successfully.")
 
@@ -94,11 +83,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 = 60
     else:
         max_input = 1024
-        max_output = 60
 
+    max_output = 60
     train_dataset = train
     eval_dataset = val
     if finetune_type == "multilingual":
@@ -124,27 +112,21 @@ def fine_tune(model_name, finetune_type, model, tokenizer, summarize_text, train
         return model_inputs
 
     tokenized_train = train_dataset.map(preprocess_function, batched=True)
-
-    # Create a small evaluation dataset
-    
     tokenized_eval = eval_dataset.map(preprocess_function, batched=True)
 
-    # Apply QLoRA only for mT5
+    # QLoRA config for mT5
     if model_name == "mT5":
-        # PEFT Configuration for Quantized Fine-tuning
         lora_config = LoraConfig(
-            r=8,  # Rank of the LoRA update matrices
-            lora_alpha=32,  # Scaling factor for the LoRA update matrices
-            lora_dropout=0.05,  # Dropout probability for the LoRA update matrices
-            bias="none",  # Whether to apply a bias to the LoRA update matrices
-            task_type=TaskType.SEQ_2_SEQ_LM  # Task type for the model
+            r=8,
+            lora_alpha=32,
+            lora_dropout=0.05,
+            bias="none",
+            task_type=TaskType.SEQ_2_SEQ_LM
         )
 
-        # Prepare model for int8 training and apply LoRA
         model = prepare_model_for_kbit_training(model)
         model = get_peft_model(model, lora_config)
     
-    # Use DataCollatorForSeq2Seq for dynamic padding
     data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model) # Initialize the DataCollatorForSeq2Seq
 
     training_args = TrainingArguments(
@@ -156,7 +138,7 @@ def fine_tune(model_name, finetune_type, model, tokenizer, summarize_text, train
         per_device_eval_batch_size=4,
         num_train_epochs=3,
         weight_decay=0.01,
-        push_to_hub=True,  # Automatically push at the end
+        push_to_hub=True,
         fp16=True,
         report_to="none",
     )
@@ -171,7 +153,6 @@ def fine_tune(model_name, finetune_type, model, tokenizer, summarize_text, train
     
     trainer.train()
     
-    # Save tokenizer and push manually
     tokenizer.save_pretrained(training_args.output_dir)
     tokenizer.push_to_hub(f"{model_name}-{finetune_type}-finetuned")