BERT_LSTM.py

import pandas as pd
import numpy as np
import torch
import sys
import torch.nn as nn

## Ensure TensorFlow is not used
import os
os.environ["USE_TF"] = "0"

## Set Random State for Reproducability
random_state = 42

# Import Hugging Face Tooling
from transformers import BertTokenizer
from transformers import BertModel
from transformers import Trainer, TrainingArguments
from transformers.modeling_outputs import SequenceClassifierOutput
import evaluate
from datasets import Dataset

# Use CPU/MPS if possible
device = None
if "google.colab" in sys.modules:
    # Running in Colab
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
else:
    # Not in Colab (e.g., Mac)
    device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")

print("Using device:", device)

train_df = pd.read_csv('../data/train_data.csv')
dev_df = pd.read_csv('../data/dev_data.csv')
test_df = pd.read_csv('../data/test_data.csv')

# Compute Class Proportions
p0 = (train_df['label'] == 0).mean() # Computes the percentage of our training dataset that has label = 0
p1 = (train_df['label'] == 1).mean() # Computes the percentage of our training dataset that has label = 1
print(f"{p0  * 100}% of our dataset has label = 0 and {p1  * 100}% of our dataset has label = 1")

# Define Custom Loss Criterion to Address Class Imbalance
class_weights = torch.tensor([p1, p0]).float().to(device)
custom_criterion = nn.CrossEntropyLoss(weight = class_weights)
print(f"Class Weights: {class_weights}")

# Fetch BERT Tokenizer from HuggingFace
bert_model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(bert_model_name)

class BertLSTMClassifier(nn.Module):
    def __init__(
        self,
        bert_model_name="bert-base-uncased",
        num_labels=2,
        hidden_dim=256,
        num_layers=2,
        bidirectional=True,
        dropout=0.3
    ):
        super().__init__()

        # Load pretrained BERT encoder
        self.bert = BertModel.from_pretrained(bert_model_name)
        bert_hidden = self.bert.config.hidden_size

        # LSTM configuration
        self.bidirectional = bidirectional
        self.lstm = nn.LSTM(
            input_size=bert_hidden,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=bidirectional
        )

        lstm_output_dim = hidden_dim * (2 if bidirectional else 1)

        # Fully connected classifier layers
        self.fc1 = nn.Linear(lstm_output_dim, 128)
        self.dropout = nn.Dropout(dropout)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, num_labels)

    def forward(self, input_ids, attention_mask, token_type_ids=None, labels=None):
        # BERT encoder
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            return_dict=True
        )

        # Extract sequence embeddings
        x = outputs.last_hidden_state  # (batch, seq_len, hidden_size)

        # Feed into LSTM
        lstm_out, (h_n, c_n) = self.lstm(x)     # (batch, seq_len, hidden_dim * num_directions)

        if self.bidirectional:
            # Take the last hidden state of LSTM for forward and backward directions
            x = torch.cat((h_n[-2], h_n[-1]), dim=1)
        else:
            # Take the last hidden state of LSTM
            x = h_n[-1]

        # Classifier layers
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        logits = self.fc2(x)

        return SequenceClassifierOutput(logits=logits)

model = BertLSTMClassifier(bert_model_name, num_labels = 2) # num_labels = 2 since we have 2 classes!

# Create `Hugging Face` Datasets [Train + Dev + Test]
train_hf_dataset = Dataset.from_pandas(train_df)
dev_hf_dataset = Dataset.from_pandas(dev_df)
test_hf_dataset = Dataset.from_pandas(test_df)

# Tokenize Text Data
def tokenize_function(row):
    tokens = tokenizer(row['text'], truncation = True, padding = 'max_length', max_length = tokenizer.model_max_length)
    row['input_ids'] = tokens['input_ids']
    row['attention_mask'] = tokens['attention_mask']
    row['token_type_ids'] = tokens['token_type_ids']
    return row

train_hf_dataset = train_hf_dataset.map(tokenize_function)
dev_hf_dataset = dev_hf_dataset.map(tokenize_function)
test_hf_dataset = test_hf_dataset.map(tokenize_function)

# Define Accuracy, Precision, Recall, and F1 Metrics from Hugging Face
accuracy_metric = evaluate.load("accuracy")
precision_metric = evaluate.load("precision")
recall_metric = evaluate.load('recall')
f1_metric = evaluate.load("f1")

# Define a compute_metrics function
def compute_metrics(eval_pred):
    # Get the model predictions
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)

    # Return Metrics
    return {
        "accuracy": accuracy_metric.compute(predictions=predictions, references=labels)['accuracy'], # Accuracy
        "pos_precision": precision_metric.compute(predictions=predictions, references=labels, pos_label = 1, average = 'binary', zero_division = 0)["precision"], # Precision on the Class w/ Label = 1 [Hate Samples]
        "pos_recall": recall_metric.compute(predictions=predictions, references=labels, pos_label = 1, average = 'binary', zero_division = 0)['recall'], # Recall on the Class w/ Label = 1 [Hate Samples]
        "pos_f1": f1_metric.compute(predictions=predictions, references=labels, pos_label = 1, average = 'binary')["f1"], # F1 Score on the Class w/ Label = 1 [Hate Samples]
        "neg_precision": precision_metric.compute(predictions=predictions, references=labels, pos_label = 0, average = 'binary', zero_division = 0)['precision'], # Precision on the Class w/ Label = 0 [Non-Hate Samples]
        "neg_recall": recall_metric.compute(predictions=predictions, references=labels, pos_label = 0, average = 'binary', zero_division = 0)['recall'], # Recall on the Class w/ Label = 0 [Non-Hate Samples]
        "neg_f1": f1_metric.compute(predictions=predictions, references=labels, pos_label = 0, average = 'binary')['f1'], # F1 Score on the Class w/ Label = 0 [Non-Hate Samples]
        "f1_macro": f1_metric.compute(predictions=predictions, references=labels, average='macro')['f1'], # Macro F1 Score
        "f1_micro": f1_metric.compute(predictions=predictions, references=labels, average='micro')['f1'], # Micro F1 Score
        "f1_weighted": f1_metric.compute(predictions=predictions, references=labels, average='weighted')['f1'], # Weighted F1 Score
    }

# Subclass the `Trainer` Class from HuggingFace to use Custom Loss Criterion
# Create a subclassed Trainer that enables us to use the custom loss function defined earlier
class SubTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs = False, **kwargs):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        loss = custom_criterion(logits, labels)
        return (loss, outputs) if return_outputs else loss

"""# **Initialize the `TrainingArguments` and `Trainer`**"""

training_args = TrainingArguments(
    output_dir="Milestone3-BERT-BiLSTM-FineTuning",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    learning_rate=5e-5,
    num_train_epochs=3,
    save_strategy="steps",      # save checkpoints every N steps
    save_steps=50,             # save every 50 steps
    eval_strategy="steps",      # evaluate every N steps
    eval_steps=50,             # evaluate every 50 steps
    logging_strategy="steps",
    logging_steps=50,          # log every 50 steps
    report_to="none",
    full_determinism=True,
    load_best_model_at_end=True, # Load best model(i.e. one that has highest dev F1 score)
    metric_for_best_model="pos_f1", # Metric used for model comparison
    greater_is_better=True # Greater F1 score is better!
)

trainer = SubTrainer(
    model=model,
    args=training_args,
    train_dataset=train_hf_dataset,
    eval_dataset=dev_hf_dataset,
    compute_metrics=compute_metrics,
)

"""# **Train the Model: `Fine-Tuning`**"""

trainer.train() # Always Resume from Last Checkpoint to Save Time
trainer.save_model('Milestone3-BERT-BiLSTM-FinalModel') # Save the Final Model
trainer.save_state() # Save the State of the Trainer (e.g. Losses, etc)

"""# **Evaluate on Train, Dev, and Test Datasets**"""

# Split: Train, Dev, or Test
def generate_evaluation_results(split):
    dataset = None
    if split == "train":
        dataset = train_hf_dataset
    elif split == "dev" or split == "validation" or split == "val":
        dataset = dev_hf_dataset
    elif split == "test":
        dataset = test_hf_dataset

    results = trainer.evaluate(eval_dataset=dataset, metric_key_prefix=split)
    df_results = pd.DataFrame([results])
    df_results.to_csv(f"BERT-BiLSTM-{split}-results.csv", index=False)
    print(f"Saved {split} evaluation metrics to BERT-BiLSTM-{split}-results.csv")

# Generate Evaluation Results on Train, Dev, and Test Splits
generate_evaluation_results("train")
generate_evaluation_results("dev")
generate_evaluation_results("test")