multitask_classifier.py

import argparse
import os
from pprint import pformat
import random
# import re
# import sys
# import time
from types import SimpleNamespace
# import pandas as pd
# from sklearn.model_selection import train_test_split
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
# from torch.nn.functional import cosine_similarity

from bert import BertModel
from my_datasets import (
    SentenceClassificationDataset,
    SentencePairDataset,
    load_multitask_data,
)
from evaluation import model_eval_multitask, test_model_multitask
from optimizer import AdamW
from sklearn.metrics import f1_score
from transformers import get_scheduler, XLNetForSequenceClassification
import datetime

TQDM_DISABLE = False

# fix the random seed
def seed_everything(seed=11711):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


BERT_HIDDEN_SIZE = 768
N_SENTIMENT_CLASSES = 5

'''
# New head for MLM task --> used for pretraining task and disabled temporarily
class BertLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transform = nn.Sequential(
            nn.Linear(config.hidden_size, config.hidden_size),
            nn.GELU(),
            nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        )
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.bias = nn.Parameter(torch.zeros(config.vocab_size))

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states) + self.bias
        return hidden_states
'''
    
class MultitaskBERT(nn.Module):
    def __init__(self, config):
        super(MultitaskBERT, self).__init__()

        self.bert = BertModel.from_pretrained(
            "bert-base-uncased", local_files_only=config.local_files_only
        )
        for param in self.bert.parameters():
            if config.option == "pretrain":
                param.requires_grad = False

            elif config.option == "finetune":
                param.requires_grad = True # Initially freeze BERT parameters
        
        self.paraphrase_head = nn.Linear(BERT_HIDDEN_SIZE * 4 + 1, 1)
        self.sentiment_layernorm = nn.LayerNorm(BERT_HIDDEN_SIZE * 2)
        self.sentiment_projection = nn.Linear(BERT_HIDDEN_SIZE * 2, N_SENTIMENT_CLASSES)
        self.sentiment_dropout = nn.Dropout(0.5)  # Stronger dropout for sentiment head
        self.paraphrase_types_head = nn.Sequential(nn.Linear(BERT_HIDDEN_SIZE, 26))
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        # self.cls = BertLMPredictionHead(self.bert.config)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        return outputs["pooler_output"]

    def etpc_forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        return outputs["last_hidden_state"][:, 0]
    
    '''
    # MLM forward pass (not used currently)
    def mlm_forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        # Pass the last hidden state of all tokens to the new cls head
        prediction_scores = self.cls(outputs["last_hidden_state"])
        return prediction_scores
    '''

    def predict_sentiment(self, input_ids, attention_mask):
        # Get both CLS token and mean of sequence
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        cls_emb = outputs["pooler_output"]
        last_hidden = outputs["last_hidden_state"]
        mean_pooling = torch.sum(last_hidden * attention_mask.unsqueeze(-1), 1) / torch.sum(attention_mask, 1, keepdim=True)
        combined = torch.cat([cls_emb, mean_pooling], dim=1)
        combined = self.sentiment_layernorm(combined)
        combined = self.sentiment_dropout(combined)
        logits = self.sentiment_projection(combined)

        return logits

    def predict_paraphrase(self, input_ids_1, attention_mask_1, input_ids_2, attention_mask_2):
        emb_1 = self.forward(input_ids_1, attention_mask_1)
        emb_2 = self.forward(input_ids_2, attention_mask_2)
        
        # Element-wise operations to capture interactions
        diff = torch.abs(emb_1 - emb_2)
        prod = emb_1 * emb_2
        cosine = F.cosine_similarity(emb_1, emb_2, dim=1).unsqueeze(1)
        
        # Combine all features
        pair_emb = torch.cat([emb_1, emb_2, diff, prod, cosine], dim=1)
        pair_emb = self.dropout(pair_emb)
        logits = self.paraphrase_head(pair_emb).squeeze(-1)
        return logits

    def predict_paraphrase_types(self, input_ids_1, attention_mask_1, input_ids_2, attention_mask_2):
        input_ids = torch.cat([input_ids_1[:, :-1], input_ids_2[:, 1:]], dim=1)
        attention_mask = torch.cat([attention_mask_1[:, :-1], attention_mask_2[:, 1:]], dim=1)
        embedding = self.etpc_forward(input_ids, attention_mask)
        embedding = self.dropout(embedding)
        logits = self.paraphrase_types_head(embedding)
        return logits


def save_model(model, optimizer, args, config, filepath):
    save_info = {
        "model": model.state_dict(),
        "optim": optimizer.state_dict(),
        "args": args,
        "model_config": config,
        "system_rng": random.getstate(),
        "numpy_rng": np.random.get_state(),
        "torch_rng": torch.random.get_rng_state(),
    }

    # Ensure parent directory exists
    parent_dir = os.path.dirname(filepath)
    if parent_dir and not os.path.exists(parent_dir):
        os.makedirs(parent_dir, exist_ok=True)

    torch.save(save_info, filepath)
    print(f"Saving the model to {filepath}.")


def f1_value_ETPC(model, dataloader, device):
    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for batch in dataloader:
            b_ids1 = batch['token_ids_1']
            b_mask1 = batch['attention_mask_1']
            b_ids2 = batch['token_ids_2']
            b_mask2 = batch['attention_mask_2']
            b_labels = batch['labels'].float().cpu().numpy()

            b_ids1 = b_ids1.to(device)
            b_mask1 = b_mask1.to(device)
            b_ids2 = b_ids2.to(device)
            b_mask2 = b_mask2.to(device)

            logits = model.predict_paraphrase_types(b_ids1, b_mask1, b_ids2, b_mask2)
            preds = (torch.sigmoid(logits) > 0.5).cpu().numpy()

            all_preds.append(preds)
            all_labels.append(b_labels)
    
    all_preds = np.vstack(all_preds)
    all_labels = np.vstack(all_labels)
    macro_f1 = f1_score(all_labels, all_preds, average='macro', zero_division=0)
    micro_f1 = f1_score(all_labels, all_preds, average='micro', zero_division=0)
    return macro_f1, micro_f1

def focal_loss(logits, targets, gamma=2.0):
    probs = torch.sigmoid(logits)
    pt = torch.where(targets == 1, probs, 1-probs)
    ce_loss = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')
    loss = (1-pt)**gamma * ce_loss
    return loss.mean()

def focal_loss(logits, targets, gamma=2.0):
    probs = torch.sigmoid(logits)
    pt = torch.where(targets == 1, probs, 1-probs)
    ce_loss = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')
    loss = (1-pt)**gamma * ce_loss
    return loss.mean()


def train_multitask(args):
    device = torch.device("cuda") if args.use_gpu else torch.device("cpu")
    # Load data
    # Create the data and its corresponding datasets and dataloader:
    sst_train_data, _, quora_train_data, sts_train_data, etpc_train_data = load_multitask_data(
        args.sst_train, 
        args.quora_train, 
        args.sts_train, 
        args.etpc_train, 
        split="train"
    )
    sst_dev_data, _, quora_dev_data, sts_dev_data, etpc_dev_data = load_multitask_data(
        args.sst_dev, 
        args.quora_dev, 
        args.sts_dev, 
        args.etpc_dev, 
        split="dev"
    )

    sst_train_dataloader = None
    sst_dev_dataloader = None
    quora_train_dataloader = None
    quora_dev_dataloader = None
    sts_train_dataloader = None
    sts_dev_dataloader = None
    etpc_train_dataloader = None
    etpc_dev_dataloader = None

    # SST dataset
    if args.task == "sst" or args.task == "multitask":
        sst_train_data = SentenceClassificationDataset(sst_train_data, args)
        sst_dev_data = SentenceClassificationDataset(sst_dev_data, args)

        sst_train_dataloader = DataLoader(
            sst_train_data,
            shuffle=True,
            batch_size=args.batch_size,
            collate_fn=sst_train_data.collate_fn,
        )
        sst_dev_dataloader = DataLoader(
            sst_dev_data,
            shuffle=False,
            batch_size=args.batch_size,
            collate_fn=sst_dev_data.collate_fn,
        )

    # Load data for the other datasets
    if args.task == "qqp" or args.task == "multitask":
        quora_train_data = SentencePairDataset(quora_train_data, args)
        quora_dev_data = SentencePairDataset(quora_dev_data, args)
        quora_train_dataloader = DataLoader(
            quora_train_data,
            shuffle=True,
            batch_size=args.batch_size,
            collate_fn=quora_train_data.collate_fn,
        )
        quora_dev_dataloader = DataLoader(
            quora_dev_data,
            shuffle=False,
            batch_size=args.batch_size,
            collate_fn=quora_dev_data.collate_fn,
        )

    if args.task == "sts" or args.task == "multitask":
        sts_train_data = SentencePairDataset(sts_train_data, args)
        sts_dev_data = SentencePairDataset(sts_dev_data, args)
        sts_train_dataloader = DataLoader(
            sts_train_data,
            shuffle=True,
            batch_size=args.batch_size,
            collate_fn=sts_train_data.collate_fn,
        )
        sts_dev_dataloader = DataLoader(
            sts_dev_data,
            shuffle=False,
            batch_size=args.batch_size,
            collate_fn=sts_dev_data.collate_fn,
        )

    if args.task == "etpc" or args.task == "multitask":
        etpc_train_data = SentencePairDataset(etpc_train_data, args)
        etpc_dev_data = SentencePairDataset(etpc_dev_data, args)

        etpc_train_dataloader = DataLoader(
            etpc_train_data,
            shuffle=True,
            batch_size=args.batch_size,
            collate_fn=etpc_train_data.collate_fn,
        )
        etpc_dev_dataloader = DataLoader(
            etpc_dev_data,
            shuffle=False,
            batch_size=args.batch_size,
            collate_fn=etpc_dev_data.collate_fn,
        )

    # Init model
    config = { 
        "hidden_dropout_prob": args.hidden_dropout_prob,
        "hidden_size": BERT_HIDDEN_SIZE,
        "data_dir": ".",
        "option": args.option,
        "local_files_only": args.local_files_only,
    }

    config = SimpleNamespace(**config)

    separator = "-" * 30
    print(separator)
    print("    BERT Model Configuration")
    print(separator)
    print(pformat({k: v for k, v in vars(args).items() if "csv" not in str(v)}))
    print(separator)

    # if args.task == "sst":
       # model = XLNetForSequenceClassification.from_pretrained('xlnet-base-cased', num_labels=5)
       # model.to(device)
    #else:
    model = MultitaskBERT(config)
    model = model.to(device)
    '''
    # Load a pre-trained model (if available)
    try:
        model.load_state_dict(torch.load("models/adapted_bert_model.bin"))
        print("Successfully loaded pre-trained model weights.")
    except FileNotFoundError:
        print("Warning: Pre-trained model not found. Starting fine-tuning with a standard BERT model.")

    '''

    weight_decay = 1e-2 if config.option == "finetune" else 0.0

    optimizer = AdamW(
        model.parameters(), # filter(lambda p: p.requires_grad, model.parameters()),
        lr=args.lr,
        weight_decay=weight_decay,
    )
    num_epochs = args.epochs

    if args.task == "sst":
        num_batches_per_epoch = len(sst_train_dataloader)
    elif args.task == "qqp":
        num_batches_per_epoch = len(quora_train_dataloader)
    elif args.task == "etpc":
        num_batches_per_epoch = len(etpc_train_dataloader)
    elif args.task == "sts":
        num_batches_per_epoch = len(sts_train_dataloader)

    num_training_steps = num_epochs * num_batches_per_epoch
    num_warmup_steps = int(0.1 * num_training_steps) # A common heuristic is 10% of total steps

    lr_scheduler = get_scheduler(
        "linear",
        optimizer=optimizer,
        num_warmup_steps=num_warmup_steps,
        num_training_steps=num_training_steps,
    )

    best_dev_acc = float("-inf")
    p_count = 0
    limit = 5 # args.epochs // 4

    for epoch in range(args.epochs):
        model.train()

        train_loss = 0
        num_batches = 0

        #sentiment classification
        
        if args.task == "sst" or args.task == "multitask":
            # Train the model on the sst dataset.
            for batch in tqdm(
                sst_train_dataloader, desc=f"train-{epoch+1:02}", disable=TQDM_DISABLE # sst_train_dataloader
            ):
                b_ids, b_mask, b_labels = (
                    batch["token_ids"],
                    batch["attention_mask"],
                    batch["labels"],
                )

                b_ids = b_ids.to(device)
                b_mask = b_mask.to(device)
                b_labels = b_labels.to(device)
                
                logits = model.predict_sentiment(b_ids, b_mask)
                loss = F.cross_entropy(logits, b_labels.view(-1)) # Third attribute label_smoothing=0.2
                
                if config.option == "finetune":
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                    lr_scheduler.step()

                train_loss += loss.item()
                num_batches += 1

        # Paraphrase similairity
        if args.task == "sts" or args.task == "multitask":
            for batch in tqdm(
                sts_train_dataloader, desc=f"train-STS-{epoch+1:02}", disable=TQDM_DISABLE
            ):
                b_ids1, b_mask1, b_ids2, b_mask2, b_labels = (
                    batch["token_ids_1"],
                    batch["attention_mask_1"],
                    batch["token_ids_2"],
                    batch["attention_mask_2"],
                    batch["labels"],
                )
                b_ids1 = b_ids1.to(device)
                b_mask1 = b_mask1.to(device)
                b_ids2 = b_ids2.to(device)
                b_mask2 = b_mask2.to(device)
                b_labels = b_labels.to(device).float()

                max_len = max(b_ids1.size(1), b_ids2.size(1))

                if b_ids1.size(1) < max_len:
                    pad_len = max_len - b_ids1.size(1)
                    b_ids1 = torch.nn.functional.pad(b_ids1, (0, pad_len), "constant", 0)
                    b_mask1 = torch.nn.functional.pad(b_mask1, (0, pad_len), "constant", 0)

                # Pad b_ids2 and b_mask2 if their length is smaller than max_len
                if b_ids2.size(1) < max_len:
                    pad_len = max_len - b_ids2.size(1)
                    b_ids2 = torch.nn.functional.pad(b_ids2, (0, pad_len), "constant", 0)
                    b_mask2 = torch.nn.functional.pad(b_mask2, (0, pad_len), "constant", 0)

                # Now, the tensors have the same sequence length and can be concatenated
                all_ids = torch.cat((b_ids1, b_ids2), dim=0)
                all_masks = torch.cat((b_mask1, b_mask2), dim=0)
                all_embs = model.forward(all_ids, all_masks)
                emb_1 = all_embs[:len(b_ids1)]
                emb_2 = all_embs[len(b_ids1):]
                cosine_scores = F.cosine_similarity(emb_1.unsqueeze(1), emb_2.unsqueeze(0), dim=2) # (batch_size, batch_size)
                line_val = 3.0 # A hyperparameter to separate positive and negative pairs (3.0 is optimum for STS-B)
                loss_labels = (b_labels > line_val).float() * 2 - 1 # Convert to -1 and 1
                negative_scores_matrix = cosine_scores.clone() # (batch_size, batch_size)
                negative_scores_matrix[loss_labels == 1] = -1.0 # Set a very low score
                hard_neg_indices = torch.argmax(negative_scores_matrix, dim=1)
                hard_negative_embs = emb_2[hard_neg_indices]
                positive_sim = F.cosine_similarity(emb_1, emb_2, dim=1) # Positive pairs
                negative_sim = F.cosine_similarity(emb_1, hard_negative_embs, dim=1) # Hard negative pairs
                positive_loss = F.mse_loss(positive_sim, torch.ones_like(positive_sim)) # Pushes the positive pairs' similarity toward 1
                negative_loss = F.mse_loss(negative_sim, -torch.ones_like(negative_sim)) # Pushes the hard negative pairs' similarity toward -1

                triplet_loss = positive_loss + negative_loss # Total triplet loss
                predicted_scores = (positive_sim + 1) * 2.5 # # Regression Loss pushes the scaled similarity score toward the true label
                regression_loss = F.mse_loss(predicted_scores, b_labels)
                alpha = 0.8  # A hyperparameter to balance the two loss types (0.6)
                loss = (alpha * triplet_loss) + ((1 - alpha) * regression_loss)

                if config.option == "finetune":
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                    lr_scheduler.step()

                train_loss += loss.item()
                num_batches += 1

        # Paraphrase detection quora
        if args.task == "qqp" or args.task == "multitask":
            # Trains the model on the qqp dataset
            for batch in tqdm(
                quora_train_dataloader, desc=f"train-QQP-{epoch+1:02}", disable=TQDM_DISABLE
            ):
                b_ids1, b_mask1, b_ids2, b_mask2, b_labels = (
                    batch["token_ids_1"].to(device),
                    batch["attention_mask_1"].to(device),
                    batch["token_ids_2"].to(device),
                    batch["attention_mask_2"].to(device),
                    batch["labels"].to(device),
                )
                b_ids1 = b_ids1.to(device)
                b_mask1 = b_mask1.to(device)
                b_ids2 = b_ids2.to(device)
                b_mask2 = b_mask2.to(device)
                b_labels = b_labels.to(device).float()

                optimizer.zero_grad()
                logits = model.predict_paraphrase(b_ids1, b_mask1, b_ids2, b_mask2)
                loss = focal_loss(logits, b_labels)

                if config.option == "finetune":
                    # Standard backward
                    loss.backward()

                    # --- Adversarial Training (FGM) ---
                    epsilon = 1e-3
                    for name, param in model.named_parameters():
                        if param.requires_grad and 'bert.embeddings' in name:
                            grad = param.grad
                            if grad is not None:
                                norm = torch.norm(grad)
                                if norm != 0:
                                    r_adv = epsilon * grad / norm
                                    param.data.add_(r_adv)
                    # Forward and backward on adversarial example
                    logits_adv = model.predict_paraphrase(b_ids1, b_mask1, b_ids2, b_mask2)
                    loss_adv = focal_loss(logits_adv, b_labels)
                    loss_adv.backward()
                    # Restore original embeddings
                    for name, param in model.named_parameters():
                        if param.requires_grad and 'bert.embeddings' in name:
                            grad = param.grad
                            if grad is not None:
                                norm = torch.norm(grad)
                                if norm != 0:
                                    r_adv = epsilon * grad / norm
                                    param.data.sub_(r_adv)

                    optimizer.step()
                    lr_scheduler.step()

                train_loss += loss.item()
                num_batches += 1

        # Paraphrase type detection
        if args.task == "etpc" or args.task == "multitask":
            # Trains the model on the etpc dataset
            for batch in tqdm(
                etpc_train_dataloader, desc=f"train-ETPC-{epoch+1:02}", disable=TQDM_DISABLE
            ):
                b_ids1, b_mask1, b_ids2, b_mask2, b_labels = (
                    batch["token_ids_1"],
                    batch["attention_mask_1"],
                    batch["token_ids_2"],
                    batch["attention_mask_2"],
                    batch["labels"],
                )
                b_ids1 = b_ids1.to(device)
                b_mask1 = b_mask1.to(device)
                b_ids2 = b_ids2.to(device)
                b_mask2 = b_mask2.to(device)
                b_labels = b_labels.to(device).float()

                optimizer.zero_grad()
                logits = model.predict_paraphrase_types(b_ids1, b_mask1, b_ids2, b_mask2)
                BCE_loss = nn.BCEWithLogitsLoss() 
                loss = BCE_loss(logits, b_labels) # F.binary_cross_entropy_with_logits
                
                if config.option == "finetune":
                    loss.backward()
                    optimizer.step()
                    lr_scheduler.step()

                train_loss += loss.item()
                num_batches += 1

        train_loss = train_loss / num_batches if num_batches > 0 else 0

        quora_train_acc, _, _, _, sst_train_acc, _, _, sts_train_corr, _, _, etpc_train_acc, _, _ = (
            model_eval_multitask(
                sst_train_dataloader, 
                quora_train_dataloader,
                sts_train_dataloader,
                etpc_train_dataloader,
                model=model,
                device=device,
                task=args.task,
            )
        )

        quora_dev_acc, _, _, _, sst_dev_acc, _, _, sts_dev_corr, _, _, etpc_dev_acc, _, _ = (
            model_eval_multitask(
                sst_dev_dataloader, 
                quora_dev_dataloader,
                sts_dev_dataloader,
                etpc_dev_dataloader,
                model=model,
                device=device,
                task=args.task,
            )
        )
        
        train_acc, dev_acc = {
            "sst": (sst_train_acc, sst_dev_acc),
            "sts": (sts_train_corr, sts_dev_corr),
            "qqp": (quora_train_acc, quora_dev_acc),
            "etpc": (etpc_train_acc, etpc_dev_acc),
        }[args.task]

        if args.task == "etpc":
            macro_f1, micro_f1 = f1_value_ETPC(model, etpc_dev_dataloader, device)
            print(f"[ETPC] Dev Macro F1: {macro_f1:.3f} -- and -- Micro F1: {micro_f1:.3f}")
        

        print(
            f"Epoch {epoch+1:02} ({args.task}): train loss :: {train_loss:.3f}, train :: {train_acc:.3f}, dev :: {dev_acc:.3f}"
        )

        if dev_acc > best_dev_acc:
            best_dev_acc = dev_acc
            p_count = 0
            save_model(model, optimizer, args, config, args.filepath)
        else:  
            p_count += 1

        if p_count >= limit:
            print(f"Early stopping triggered. No improvement for {limit} epochs.")
            break # Exit the training loop
        

def test_model(args):
    with torch.no_grad():
        device = torch.device("cuda") if args.use_gpu else torch.device("cpu")
        saved = torch.load(args.filepath, weights_only=False) # Temporarily set weights_only=False
        config = saved["model_config"]

        model = MultitaskBERT(config)
        model.load_state_dict(saved["model"])
        model = model.to(device)
        print(f"Loaded model to test from {args.filepath}")

        return test_model_multitask(args, model, device)


def get_args():
    parser = argparse.ArgumentParser()

    # Training task
    parser.add_argument(
        "--task",
        type=str,
        help='choose between "sst","sts","qqp","etpc","multitask" to train for different tasks ',
        choices=("sst", "sts", "qqp", "etpc", "multitask"),
        default="sst",
    )

    # Model configuration
    parser.add_argument("--seed", type=int, default=11711)
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument(
        "--option",
        type=str,
        help="pretrain: the BERT parameters are frozen; finetune: BERT parameters are updated",
        choices=("pretrain", "finetune"),
        default="pretrain",
    )
    parser.add_argument("--use_gpu", action="store_true")

    args, _ = parser.parse_known_args()

    # Dataset paths
    parser.add_argument("--sst_train", type=str, default="data/sst-sentiment-train.csv")
    parser.add_argument("--sst_dev", type=str, default="data/sst-sentiment-dev.csv")
    parser.add_argument("--sst_test", type=str, default="data/sst-sentiment-test-student.csv")

    parser.add_argument("--quora_train", type=str, default="data/quora-paraphrase-train.csv")
    parser.add_argument("--quora_dev", type=str, default="data/quora-paraphrase-dev.csv")
    parser.add_argument("--quora_test", type=str, default="data/quora-paraphrase-test-student.csv")

    parser.add_argument("--sts_train", type=str, default="data/sts-similarity-train.csv")
    parser.add_argument("--sts_dev", type=str, default="data/sts-similarity-dev.csv")
    parser.add_argument("--sts_test", type=str, default="data/sts-similarity-test-student.csv")

    # TODO
    # You should split the train data into a train and dev set first and change the
    # default path of the --etpc_dev argument to your dev set.

    parser.add_argument("--etpc_train", type=str, default="data/etpc-paraphrase-train_split.csv")
    parser.add_argument("--etpc_dev", type=str, default="data/etpc-paraphrase-dev.csv")
    parser.add_argument("--etpc_test", type=str, default="data/etpc-paraphrase-detection-test-student.csv")

    # Output paths
    parser.add_argument(
        "--sst_dev_out",
        type=str,
        default=(
            "predictions/bert/sst-sentiment-dev-output.csv"
            if not args.task == "multitask"
            else "predictions/bert/multitask/sst-sentiment-dev-output.csv"
        ),
    )
    parser.add_argument(
        "--sst_test_out",
        type=str,
        default=(
            "predictions/bert/sst-sentiment-test-output.csv"
            if not args.task == "multitask"
            else "predictions/bert/multitask/sst-sentiment-test-output.csv"
        ),
    )

    parser.add_argument(
        "--quora_dev_out",
        type=str,
        default=(
            "predictions/bert/quora-paraphrase-dev-output.csv"
            if not args.task == "multitask"
            else "predictions/bert/multitask/quora-paraphrase-dev-output.csv"
        ),
    )
    parser.add_argument(
        "--quora_test_out",
        type=str,
        default=(
            "predictions/bert/quora-paraphrase-test-output.csv"
            if not args.task == "multitask"
            else "predictions/bert/multitask/quora-paraphrase-test-output.csv"
        ),
    )

    parser.add_argument(
        "--sts_dev_out",
        type=str,
        default=(
            "predictions/bert/sts-similarity-dev-output.csv"
            if not args.task == "multitask"
            else "predictions/bert/multitask/sts-similarity-dev-output.csv"
        ),
    )
    parser.add_argument(
        "--sts_test_out",
        type=str,
        default=(
            "predictions/bert/sts-similarity-test-output.csv"
            if not args.task == "multitask"
            else "predictions/bert/multitask/sts-similarity-test-output.csv"
        ),
    )

    parser.add_argument(
        "--etpc_dev_out",
        type=str,
        default=(
            "predictions/bert/etpc-paraphrase-detection-dev-output.csv"
            if not args.task == "multitask"
            else "predictions/bert/multitask/etpc-paraphrase-detection-dev-output.csv"
        ),
    )
    parser.add_argument(
        "--etpc_test_out",
        type=str,
        default=(
            "predictions/bert/etpc-paraphrase-detection-test-output.csv"
            if not args.task == "multitask"
            else "predictions/bert/multitask/etpc-paraphrase-detection-test-output.csv"
        ),
    )

    # Hyperparameters
    parser.add_argument("--batch_size", help="sst: 64 can fit a 12GB GPU", type=int, default=16) # 64
    parser.add_argument("--hidden_dropout_prob", type=float, default=0.3)
    parser.add_argument(
        "--lr",
        type=float,
        help="learning rate, default lr for 'pretrain': 1e-3, 'finetune': 1e-5",
        default=1e-3 if args.option == "pretrain" else 1e-5,
    )
    parser.add_argument("--local_files_only", action="store_true")

    args = parser.parse_args()
    return args

'''
def pre_traintask(args):
    """
    Performs additional pre-training of the BERT model on a domain-specific dataset.
    """
    device = torch.device("cuda") if args.use_gpu else torch.device("cpu")

    model_save_path = "models/adapted_bert_model.bin"
    
    # Check if the adapted model already exists
    if os.path.exists(model_save_path):
        print(f"Skipping pre-training. Adapted model '{model_save_path}' already exists.")
        return # Exit the function early
    
    print("--- Step 1: Additional Pre-training on IMDb Data ---")

    # 1. Prepare the data for MLM
    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
    imdb_dataset = load_dataset("imdb")
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)
    
    def tokenize_function(examples):
        return tokenizer(examples['text'], truncation=True, max_length=512)
    
    tokenized_imdb = imdb_dataset.map(tokenize_function, batched=True, remove_columns=["text", "label"])
    mlm_dataloader = DataLoader(
        tokenized_imdb['train'],
        batch_size=args.batch_size,
        collate_fn=data_collator,
    )

    # 2. Initialize the model and optimizer for pre-training
    config = {
        "hidden_dropout_prob": args.hidden_dropout_prob,
        "hidden_size": BERT_HIDDEN_SIZE,
        "data_dir": ".",
        "option": args.option,
        "local_files_only": args.local_files_only,

    }

    config = SimpleNamespace(**config)
    model = MultitaskBERT(config)
    model.to(device)

    # Unfreeze all BERT and MLM head parameters for pre-training
    for param in model.parameters():
        param.requires_grad = True

    optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=0.0)
    pre_train_epochs = 5 # Fewer epochs for pre-training
    # 3. The pre-training loop
    model.train()

    for epoch in range(pre_train_epochs): # Fewer epochs for pre-training
        print(f"Starting MLM Pre-training Epoch {epoch + 1}")
        total_loss = 0
        for batch in tqdm(mlm_dataloader, desc="MLM Pre-training", disable=TQDM_DISABLE):
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)
            
            optimizer.zero_grad()
            p_scores = model.mlm_forward(input_ids, attention_mask)
            loss = F.cross_entropy(p_scores.view(-1, tokenizer.vocab_size), labels.view(-1), ignore_index=-100)
            
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        avg_loss = total_loss / len(mlm_dataloader)
        print(f"Epoch {epoch + 1} with MLM Loss: {avg_loss:.4f}")

    # Save the domain-adapted BERT model
    torch.save(model.state_dict(), "models/adapted_bert_model.bin")
    print("Domain-adapted model saved.")
'''

if __name__ == "__main__":
    args = get_args()
    args.filepath = f"models/{args.option}-{args.epochs}-{args.lr}-{args.task}.pt"  # save path
    seed_everything(args.seed)  # fix the seed for reproducibility
    # pre_traintask(args) # Uncomment to run pre-training if needed
    train_multitask(args)
    test_model(args)