From 96a0c45e84d86c8902acf1fc94a08b1d8bebd78a Mon Sep 17 00:00:00 2001 From: charlie-rasberry Date: Thu, 26 Feb 2026 18:21:13 +0000 Subject: [PATCH] Added implementation for single task roberta, using args for everything made it simple --- src/__pycache__/model.cpython-313.pyc | Bin 2924 -> 4153 bytes src/model.py | 16 ++++++++- src/train.py | 46 ++++++++++++++++++++------ 3 files changed, 51 insertions(+), 11 deletions(-) diff --git a/src/__pycache__/model.cpython-313.pyc b/src/__pycache__/model.cpython-313.pyc index 2f87f939135b2fd3b053114ab963abea1927be0b..a2ed4743633a6195e96fdba4eca183bddd147b16 100644 GIT binary patch delta 1720 zcmZuxO-vg{6rNe{uK(~K!GOVlV-T<{At5Slk_ISgL<=-jwIPa>m8`YBfD>a!vqr7j z1Dp_tM#79*ps7@j;Z`-5q^f%8rRk*taR^$`Q%^{(6bYn6?WOOHF+@>E+Hc;@d-L9# z@4cD7>Yv2?hd!T&z<9E}nftY->~Ewi<*tuCjU-NzN)t({<(qB&F5R7UUb#hxu|!Cc zBwgd4F)#DRPxZFKIV%NohG9Eo)lj2_EYpu2n*<5a9N`V`laI)kM4^T z-IMes{g)*c)I;!B3tzChr#iDUx+_llE87UG)9b+MOfsOr3d}*APrQcpaN^vtDdB^( z#n-?Z^$6@}&>ItP@mZ?T5MQPDhXTo5VJ2@*8rIz##jKe>@WZx@T-^&Z$`1Gr{cwsw z`W3l-qMVfZq%eL;GsVJmZpQY`=CWC{kg{^0o3=b_8d=+20y0tsW6snRTei&n zG^>Z;wzn`pm&)V~%Q7vy2K|BK6f;Vu?M_t`sg#9vh{0HN*zx!IksR;9%wA~wFTMM9^=qT6qwCr4@*DYQt$jPKea~74 zc3KA>f3aIP4i@gcNc_J45Im3!mZ&@2{ZxqmaSF< zTg?||awW@BL5ozyVysyhP8Zpy2FuzFuah-hTo4+JVod-nlD*pcwVra%!^ld{c5PzW zxfiTk^On8qA3Tfp??n55Fdm0~9R7Z0Cpx$r99mZQLgBSYIkNutZYU0aVC*mZH`VPx z*S4#RwSY@a7Kf`G^W->Nbmp5481lw%LtFf3*E=eflXdccUE1&#Wlfk7UhntB85mM8 z0NJvzyI(Gm7m}j1zo?Ow?&kzRgcC-5o8F|&kGiDa=_d@IpaXmt4v?4wESVNycn6q> ziJirK2%-{aP^v?~7GgmF+j-oT3CAAWFx|N!)~S_TmX2rAQy2Fp_vc6n^pA2 zF7mWnZ^gG%qqZk)lrpp80XS`1t(ot;yRRX$whBqj=gmrKg}XO!*_k!qp)gDAL;MSt ziYznRJ73D>tzOiyQCMa^KJ4kE?cDO5y)HtFMt5DE6D|jt4o28L)-M3h9hA}+9zrMR f4!Q9M2_FT$w0kAJmMiCu2#7~b{HbSCMo0ev-12$U delta 807 zcmYLHO>fgc5Z!eWJ9hk$wrLs?nx>(Gno1?DsGxil4xG{=r22xyO2#BkYbnHN>A9E)=8TJP2>D}p{;}7AdGHSc`;LZ z;;1WB^dkgJrZtJ9F>^&4WmKj&2~(YdqpwO}u_CfD{H)_M$2*KOrAbzX#|TWY2~>?c zMh-K$n_~ahwZuxB3nMAtqIIJTQ>=mq%4}+LnZKkqwfI;1!pPkWyMF(`Ky63jccs$4 zl>bcm^UUkr&0c;qM2V2~&`Abb$+#EH@! z1@cBP@IBypZLO7LBfsBE^ij^ZT&Jo%vDa8Q4bXXCj?Xl)DUbHPKx-HfH^)R zPx5d2LTesrq#~f_aKQ|+MCpa_(1UJb46R4B&a1|o)+{P)`G4=BU+O%DrPSG`TkqeA z_2Zh&f65nRaS=Gj=M{U}rcg%~i5GD_^#esz1sMFdQImxOD7O6}#vydJ7kA>IA1#UstYR99{ExXr l&+?WvTTfln>*ufBPrI$b28zU-(H;-2Now-f*2k<6{snFMwKV_$ diff --git a/src/model.py b/src/model.py index 0996802..ba796c1 100644 --- a/src/model.py +++ b/src/model.py @@ -10,7 +10,21 @@ import torch.nn as nn # Each nn.linear is used to map RoBERTa's hidden representation onto the output space of each task head # Each hidden representation is size 768 -class Model(nn.Module): + +class SingleTaskModel(nn.Module): # SINGLE TASK MODEL ARCHITECTURE + def __init__(self, task_name, num_classes, dropout_rate=0.2): + super().__init__() + self.encoder = XLMRobertaModel.from_pretrained("FacebookAI/xlm-roberta-base") + self.droput = nn.Dropout(dropout_rate) + self.head = nn.Linear(self.encoder.config.hidden_size, num_classes) + self.task_name = task_name + def forward(self, input_ids, attention_mask): + outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask) + output= self.droput(outputs.last_hidden_state[:, 0, :]) + logits = self.head(output) + return {self.task_name: logits} + +class Model(nn.Module): # MULTITASK MODEL ARCHITECTURE def __init__(self, dropout_rate=0.2): # Try other p values super().__init__() self.encoder = XLMRobertaModel.from_pretrained("FacebookAI/xlm-roberta-base") diff --git a/src/train.py b/src/train.py index c2e289e..69a507f 100644 --- a/src/train.py +++ b/src/train.py @@ -18,7 +18,14 @@ from sklearn.utils.class_weight import compute_class_weight from dataset import ReviewDataset -from model import Model +from model import Model, SingleTaskModel + + + + +# ======================================================================= +# Multitask implementation +# ======================================================================= # NFR5, reproducibility SEED = 4321 @@ -41,6 +48,8 @@ def compute_weights(df, column, device): # python src/train.py --epochs 15 NOTE: 8 - 12 epochs has seen best results so far def parse_args(): parser = argparse.ArgumentParser(description="RECLASS, Multitask learning for review classification.") + parser.add_argument("--mode", type=str, default="mtl", choices=["mtl", "stl"], help="Choose between 'mtl' (multitask learning) and 'stl' (single task learning).") + parser.add_argument("--task", type=str, default="all", choices=["all", "bug_report", "feature_request", "aspect", "aspect_sentiment"], help="Specific task to train for stl usage only" ) parser.add_argument("--dataset", type=str, default="original", choices=["original", "boosted"], help="Choose between 'original' and 'boosted' dataset.") parser.add_argument("--batch_size", type=int, default=16, help="Keep to 16 or 8 for 8GB VRAM") parser.add_argument("--epochs", type=int, default=5, help="Maxiumum training epochs.") @@ -81,7 +90,24 @@ def main(): validation_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False) # FR3, shared multilingual model with task-specific heads - model = Model().to(device) + if args.mode == "mtl": + model = Model().to(device) + active_tasks = ['bug_report', 'feature_request', 'aspect', 'aspect_sentiment'] + run_name = f"mtl_{args.dataset}" + else: + if args.task == "all": + raise ValueError("For single task learning, please specify a task using --task argument.") + + task_classes = { + 'bug_report': 2, + 'feature_request': 2, + 'aspect': 6, + 'aspect_sentiment': 3 + } + model = SingleTaskModel(args.task, task_classes[args.task]).to(device) + active_tasks = [args.task] + run_name = f"stl_{args.task}_{args.dataset}" + train_df = pd.read_csv(train) # Class weights @@ -128,7 +154,7 @@ def main(): # ------------------- Training loop ------------------- timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') - writer = SummaryWriter('runs/fashion_trainer_{}'.format(timestamp)) + writer = SummaryWriter(f'runs/reclass_{run_name}_{timestamp}') best_f1 = 0.0 patience_counter = 0 @@ -153,7 +179,7 @@ def main(): outputs = model(input_ids, attention_mask) loss = 0 - for task in criterions.keys(): + for task in active_tasks: labels = batch[task].to(device) loss += criterions[task](outputs[task], labels) @@ -177,8 +203,8 @@ def main(): model.eval() total_val_loss = 0.0 - all_preds = {task: [] for task in criterions.keys()} - all_labels ={task: [] for task in criterions.keys()} + all_preds = {task: [] for task in active_tasks} + all_labels ={task: [] for task in active_tasks} with torch.no_grad(): for batch in validation_loader: @@ -188,7 +214,7 @@ def main(): outputs = model(input_ids, attention_mask) v_loss = 0.0 # batch validation loss - for task in criterions.keys(): + for task in active_tasks: labels = batch[task].to(device) v_loss += criterions[task](outputs[task], labels).item() # detatch .item(*) @@ -203,8 +229,8 @@ def main(): # FR11, Performance evaluation print("\nValidation Metrics (MACRO F1):") epoch_f1 = [] - for task in criterions.keys(): - task_f1 = f1_score(all_labels[task], all_preds[task], average='macro') + for task in active_tasks: + task_f1 = f1_score(all_labels[task], all_preds[task], average='macro', zero_division=0) epoch_f1.append(task_f1) writer.add_scalar(f"F1/val_{task}", task_f1, epoch) print(f" {task}: {task_f1:.4f}") @@ -218,7 +244,7 @@ def main(): best_f1 = avg_macro_f1 patience_counter = 0 # Save the model with a name for the type of dataset and epoch for later analysis - model_save_path = f"outputs/best_model_{args.dataset}.pt" + model_save_path = f"outputs/best_model_{run_name}.pt" torch.save(model.state_dict(), model_save_path) print(" New best model saved to:", model_save_path) else: