File size: 4,750 Bytes
c6f9ba7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import torch
from transformers import MBart50Tokenizer, MBartForConditionalGeneration # type: ignore
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, TaskType
from dotenv import load_dotenv
import wandb
import json
from utils.helper import TextPreprocessor
from utils.trainer import train_model
load_dotenv()
class MBart50Finetuner:
"""Class to handle fine-tuning of mBART50 model for translation tasks."""
def __init__(self, config_path="config.json"):
"""Initialize with configuration file."""
with open(config_path, "r") as json_file:
cfg = json.load(json_file)
self.args = cfg["mbart50"]["args"]
self.lora_config = cfg["mbart50"]["lora_config"]
# Constants
self.max_len = self.args["max_len"]
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.id = self.args["id"]
self.initial_learning_rate = self.args["initial_learning_rate"]
self.model_name = self.args["model_name"]
self.src_lang = self.args["src_lang"]
self.tgt_lang = self.args["tgt_lang"]
self.wandb_project = self.args["wandb_project"]
self.output_dir = self.args["output_dir"]
self.name = "mbart50"
self.model = None
self.tokenizer = None
self.train_dataset = None
self.val_dataset = None
self.test_dataset = None
def setup_wandb(self):
"""Initialize Weights & Biases for experiment tracking."""
wandb.login(key=os.environ.get("WANDB_API"), relogin=True)
wandb.init(project=self.wandb_project, name="mbart50-finetune-lora")
def load_model_and_tokenizer(self):
"""Load the mBART model and tokenizer."""
self.tokenizer = MBart50Tokenizer.from_pretrained(self.model_name)
self.model = MBartForConditionalGeneration.from_pretrained(self.model_name)
self.tokenizer.src_lang = self.src_lang
self.tokenizer.tgt_lang = self.tgt_lang
def load_datasets(self):
"""Load training, validation, and test datasets."""
data_files = {
"train": "data/train_cleaned_dataset.csv",
"test": "data/test_cleaned_dataset.csv",
"val": "data/val_cleaned_dataset.csv",
}
if self.id is not None:
training_parts = [
f"[{(i * 200000) + 1 if i > 0 else ''}:{(i + 1) * 200000 if i < 10 else ''}]"
for i in range(11)
]
self.train_dataset = load_dataset(
"csv", data_files=data_files, split=f"train{training_parts[self.id]}"
)
self.test_dataset = load_dataset("csv", data_files=data_files, split="test")
self.val_dataset = load_dataset(
"csv", data_files=data_files, split="val[:20000]"
)
else:
self.train_dataset = load_dataset(
"csv", data_files=data_files, split="train[:1000000]"
)
self.test_dataset = load_dataset("csv", data_files=data_files, split="test[:100000]")
self.val_dataset = load_dataset("csv", data_files=data_files, split="val[:100000]")
def configure_lora(self):
"""Apply LoRA configuration to the model."""
lora_config = LoraConfig(
task_type=TaskType.SEQ_2_SEQ_LM,
r=self.lora_config["r"],
lora_alpha=self.lora_config["lora_alpha"],
target_modules=self.lora_config["target_modules"],
lora_dropout=self.lora_config["lora_dropout"],
)
self.model = get_peft_model(self.model, lora_config) # type: ignore
def finetune(self):
"""Orchestrate the fine-tuning process."""
self.setup_wandb()
self.load_model_and_tokenizer()
self.load_datasets()
preprocessor = TextPreprocessor(self.tokenizer, self.max_len, name="mbart50")
tokenized_train_dataset = preprocessor.preprocess_dataset(self.train_dataset)
tokenized_eval_dataset = preprocessor.preprocess_dataset(self.val_dataset)
self.configure_lora()
self.model.print_trainable_parameters() # type: ignore
train_model(
model=self.model,
tokenizer=self.tokenizer,
train_dataset=tokenized_train_dataset,
eval_dataset=tokenized_eval_dataset,
output_dir=self.output_dir,
initial_learning_rate=self.initial_learning_rate,
name=self.name,
val_dataset=self.val_dataset,
)
if __name__ == "__main__":
finetuner = MBart50Finetuner()
finetuner.finetune()
|