tweetSADataset.py

import os
import random
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler

#####################################
# Create a Dataset Class
#####################################
class TweetSADataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
        
        # label mapping: already 0 = positive, 1 = negative
        self.label_to_id = {0: "positive", 1: "negative"}
        
        # pad token info
        self.pad_token_id = tokenizer.pad_token_id

        self.samples = []
        for text, label in zip(self.texts, self.labels):
            # Tokenize using RoBERTa tokenizer
            inputs = self.tokenizer(
                text, 
                add_special_tokens=True, 
                truncation=True, 
                max_length=self.max_length
            )
            input_ids = inputs["input_ids"]
            self.samples.append({"ids": input_ids, "label": label})

    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        return self.samples[idx]

    def padding(self, inputs, max_length=-1):
        if max_length < 0:
            max_length = max(len(x) for x in inputs)
        padded_inputs = []
        for seq in inputs:
            if len(seq) < max_length:
                seq = seq + [self.pad_token_id]*(max_length - len(seq))
            else:
                seq = seq[:max_length]
            padded_inputs.append(seq)
        return padded_inputs

    def collate_fn(self, batch):
        batch_ids = [x["ids"] for x in batch]
        labels = [x["label"] for x in batch]

        pad_batch_ids = self.padding(batch_ids)
        tensor_batch_ids = torch.tensor(pad_batch_ids, dtype=torch.long)
        tensor_labels = torch.tensor(labels, dtype=torch.long)
        return tensor_batch_ids, tensor_labels
    




#####################################
# Define Training and Evaluation Functions
#####################################
def compute_metrics(preds, labels):
    preds = np.array(preds)
    labels = np.array(labels)
    TP = np.sum((preds == 0) & (labels == 0))
    FN = np.sum((preds == 1) & (labels == 0))
    FP = np.sum((preds == 0) & (labels == 1))
    TN = np.sum((preds == 1) & (labels == 1))
    # F1 for positive (label=0)
    precision_pos = TP/(TP+FP) if TP+FP>0 else 0
    recall_pos = TP/(TP+FN) if TP+FN>0 else 0
    f1_pos = 2*precision_pos*recall_pos/(precision_pos+recall_pos) if precision_pos+recall_pos>0 else 0

    # F1 for negative (label=1)
    precision_neg = TN/(TN+FN) if TN+FN>0 else 0
    recall_neg = TN/(TN+FP) if TN+FP>0 else 0
    f1_neg = 2*precision_neg*recall_neg/(precision_neg+recall_neg) if precision_neg+recall_neg>0 else 0

    macro_f1 = (f1_pos+f1_neg)/2
    return macro_f1, (TP, FN, FP, TN), f1_pos, f1_neg