ConstraintClassifier.py

"""ConstraintClassifier.ipynb

Automatically generated by Colab.

Original file is located at
https://colab.research.google.com/drive/1OK-oih_tx6s69WbdNNOn5gcP0EdVQF2x?authuser=1
"""

import os
import torch
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from transformers import RobertaTokenizer, RobertaModel

#############################################
# Hyperparameters and File Paths (edit as needed)
EPOCHS_PHASE1 = 10         # # of epochs to train the two initial classifiers
EPOCHS_PHASE3 = 5          # # of epochs to train the third classifier
BATCH_SIZE = 16
LEARNING_RATE = 2e-5
MAX_LENGTH = 128
# "Constraint" has 2 classes: 0 (soft) and 1 (hard).
NUM_CONSTRAINT_CLASSES = 2

# Map each possible 'type' to an integer ID:
TYPE2ID = {
    'language_requirement': 0,
    'system': 1,
    'prerequisite_knowledge': 2,
    'accessibility': 3,
    'budget': 4,
    'learning_style': 5,
    'time_commitment': 6,
    'level_of_depth': 7,
    'preferred_topics': 8,
    'format_preferences': 9
}
NUM_TYPE_CLASSES = len(TYPE2ID)

LABELLED_FILE = r"/content/Combined_Dataset_Constraint.csv" # Change this
UNLABELLED_FILE = r"/content/Unlabelled_Constraint_Data.csv" # Change this
MODEL_SAVE_PATH = "model3_weights.pth"
USE_AMP = True

#############################################

if torch.cuda.is_available():
    torch.backends.cudnn.benchmark = True

# Dataset now can return two labels: constraint_label and type_label
class ConstraintDataset(Dataset):
    def __init__(self, texts, constraints=None, types=None,
                 tokenizer=None, max_length=128):
        self.texts = texts
        self.constraints = constraints  # e.g. [0 or 1]
        self.types = types             # e.g. [0..9] or None
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = str(self.texts[idx])
        inputs = self.tokenizer(
            text,
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors="pt"
        )
        input_ids = inputs["input_ids"].squeeze(0)
        attention_mask = inputs["attention_mask"].squeeze(0)

        # If constraints or types aren't provided, we return just the text.
        if self.constraints is not None and self.types is not None:
            c_label = torch.tensor(self.constraints[idx], dtype=torch.long)
            t_label = torch.tensor(self.types[idx], dtype=torch.long)
            return input_ids, attention_mask, c_label, t_label
        else:
            return input_ids, attention_mask

# Multi-task model: one head for constraint, one for type
class RobertaMultiTaskClassifier(nn.Module):
    def __init__(self, num_constraint_classes, num_type_classes):
        super(RobertaMultiTaskClassifier, self).__init__()
        self.roberta = RobertaModel.from_pretrained('roberta-base')
        self.dropout = nn.Dropout(0.1)
        self.constraint_classifier = nn.Linear(
            self.roberta.config.hidden_size, num_constraint_classes
        )
        self.type_classifier = nn.Linear(
            self.roberta.config.hidden_size, num_type_classes
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.last_hidden_state[:, 0]
        pooled_output = self.dropout(pooled_output)

        constraint_logits = self.constraint_classifier(pooled_output)
        type_logits = self.type_classifier(pooled_output)
        return constraint_logits, type_logits

# Training for multi-task: we compute constraint loss + type loss
def train_epoch(model, data_loader, optimizer, device, criterion, use_amp):
    model.train()
    losses = []
    total_constraint_correct = 0
    total_type_correct = 0
    total_examples = 0

    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

    for batch in data_loader:
        input_ids, attention_mask, c_labels, t_labels = [x.to(device) for x in batch]
        optimizer.zero_grad()

        with torch.cuda.amp.autocast(enabled=use_amp):
            c_logits, t_logits = model(input_ids, attention_mask)
            loss_c = criterion(c_logits, c_labels)
            loss_t = criterion(t_logits, t_labels)
            loss = loss_c + loss_t

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        losses.append(loss.item())

        c_preds = torch.argmax(c_logits, dim=1)
        t_preds = torch.argmax(t_logits, dim=1)
        total_constraint_correct += (c_preds == c_labels).sum().item()
        total_type_correct += (t_preds == t_labels).sum().item()
        total_examples += c_labels.size(0)

    avg_loss = np.mean(losses)
    avg_acc_constraint = total_constraint_correct / total_examples
    avg_acc_type = total_type_correct / total_examples
    return avg_loss, avg_acc_constraint, avg_acc_type

# Evaluation for multi-task
def eval_model(model, data_loader, device, criterion, use_amp):
    model.eval()
    losses = []
    total_constraint_correct = 0
    total_type_correct = 0
    total_examples = 0

    with torch.no_grad():
        for batch in data_loader:
            input_ids, attention_mask, c_labels, t_labels = [x.to(device) for x in batch]
            with torch.cuda.amp.autocast(enabled=use_amp):
                c_logits, t_logits = model(input_ids, attention_mask)
                loss_c = criterion(c_logits, c_labels)
                loss_t = criterion(t_logits, t_labels)
                loss = loss_c + loss_t
            losses.append(loss.item())

            c_preds = torch.argmax(c_logits, dim=1)
            t_preds = torch.argmax(t_logits, dim=1)
            total_constraint_correct += (c_preds == c_labels).sum().item()
            total_type_correct += (t_preds == t_labels).sum().item()
            total_examples += c_labels.size(0)

    avg_loss = np.mean(losses)
    avg_acc_constraint = total_constraint_correct / total_examples
    avg_acc_type = total_type_correct / total_examples
    return avg_loss, avg_acc_constraint, avg_acc_type

# For pseudo-labeling, we only need the constraint predictions
def predict_constraint(model, data_loader, device, use_amp):
    model.eval()
    preds = []
    with torch.no_grad():
        for batch in data_loader:
            # unlabeled data won't have c_labels, t_labels
            if len(batch) == 4:
                input_ids, attention_mask, _, _ = [x.to(device) for x in batch]
            else:
                input_ids, attention_mask = [x.to(device) for x in batch]
            with torch.cuda.amp.autocast(enabled=use_amp):
                c_logits, _ = model(input_ids, attention_mask)
            pred_c = torch.argmax(c_logits, dim=1)
            preds.extend(pred_c.cpu().numpy())
    return preds

# Partial supervision: type is only known for originally labeled data
# For pseudo-labeled data, we have no type => store -1 => skip in type loss
def train_epoch_partial(model, data_loader, optimizer, device, criterion, use_amp):
    model.train()
    losses = []
    total_constraint_correct = 0
    total_type_correct = 0
    total_constraint_examples = 0
    total_type_examples = 0

    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

    for batch in data_loader:
        input_ids, attention_mask, c_labels, t_labels = [x.to(device) for x in batch]
        optimizer.zero_grad()

        with torch.cuda.amp.autocast(enabled=use_amp):
            c_logits, t_logits = model(input_ids, attention_mask)
            # constraint loss on all examples
            loss_c = criterion(c_logits, c_labels)

            # type loss only where type != -1
            mask = (t_labels != -1)
            if mask.any():
                valid_indices = torch.where(mask)[0]
                loss_t = criterion(t_logits[valid_indices], t_labels[valid_indices])
                loss = loss_c + loss_t
            else:
                loss = loss_c

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        losses.append(loss.item())

        c_preds = torch.argmax(c_logits, dim=1)
        total_constraint_correct += (c_preds == c_labels).sum().item()
        total_constraint_examples += c_labels.size(0)

        # Evaluate type accuracy only for labeled subset
        if mask.any():
            valid_indices = torch.where(mask)[0]
            t_preds = torch.argmax(t_logits[valid_indices], dim=1)
            total_type_correct += (t_preds == t_labels[valid_indices]).sum().item()
            total_type_examples += len(valid_indices)

    avg_loss = np.mean(losses)
    avg_acc_constraint = (total_constraint_correct / total_constraint_examples
                          if total_constraint_examples else 0)
    avg_acc_type = (total_type_correct / total_type_examples
                    if total_type_examples else 0)
    return avg_loss, avg_acc_constraint, avg_acc_type

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
    pin_memory = True if device.type == "cuda" else False

    ###################################
    # Load labeled data
    ###################################
    df = pd.read_csv(LABELLED_FILE)
    # Convert 'Constraint' from {1,2} to {0,1}
    df['Constraint'] = pd.to_numeric(df['Constraint'], errors='coerce')
    df.dropna(subset=['Constraint'], inplace=True)
    df['Constraint'] = df['Constraint'].astype(int) - 1

    # Convert 'type' strings to IDs. Drop examples that don't match any known type.
    df.dropna(subset=['Type'], inplace=True)
    df['type_id'] = df['Type'].apply(lambda t: TYPE2ID.get(str(t), np.nan))
    df.dropna(subset=['type_id'], inplace=True)

    # Prepare lists
    texts = df["Text"].tolist()
    constraints = df["Constraint"].tolist()               # [0 or 1]
    types_ = df["type_id"].astype(int).tolist()           # [0..9]

    # Phase 1/2/3 splits
    X_temp, X_test, c_temp, c_test, t_temp, t_test = train_test_split(
        texts, constraints, types_, test_size=0.1, random_state=42, stratify=constraints
    )
    X_train, X_val, c_train, c_val, t_train, t_val = train_test_split(
        X_temp, c_temp, t_temp,
        test_size=0.2222, random_state=42, stratify=c_temp
    )

    # Build datasets
    train_dataset = ConstraintDataset(X_train, c_train, t_train, tokenizer, MAX_LENGTH)
    val_dataset   = ConstraintDataset(X_val,   c_val,   t_val,   tokenizer, MAX_LENGTH)
    test_dataset  = ConstraintDataset(X_test,  c_test,  t_test,  tokenizer, MAX_LENGTH)

    train_loader = DataLoader(train_dataset, BATCH_SIZE, shuffle=True, pin_memory=pin_memory)
    val_loader   = DataLoader(val_dataset,   BATCH_SIZE, pin_memory=pin_memory)
    test_loader  = DataLoader(test_dataset,  BATCH_SIZE, pin_memory=pin_memory)

    ###################################
    # Phase 1: Train two classifiers on labeled data
    ###################################
    print("\nPhase 1: Training two classifiers on labeled data")
    model1 = RobertaMultiTaskClassifier(NUM_CONSTRAINT_CLASSES, NUM_TYPE_CLASSES).to(device)
    model2 = RobertaMultiTaskClassifier(NUM_CONSTRAINT_CLASSES, NUM_TYPE_CLASSES).to(device)

    optimizer1 = optim.AdamW(model1.parameters(), lr=LEARNING_RATE)
    optimizer2 = optim.AdamW(model2.parameters(), lr=LEARNING_RATE)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(1, EPOCHS_PHASE1 + 1):
        # Train each model
        train_loss1, train_c_acc1, train_t_acc1 = train_epoch(
            model1, train_loader, optimizer1, device, criterion, USE_AMP
        )
        val_loss1, val_c_acc1, val_t_acc1 = eval_model(
            model1, val_loader, device, criterion, USE_AMP
        )

        train_loss2, train_c_acc2, train_t_acc2 = train_epoch(
            model2, train_loader, optimizer2, device, criterion, USE_AMP
        )
        val_loss2, val_c_acc2, val_t_acc2 = eval_model(
            model2, val_loader, device, criterion, USE_AMP
        )

        print(f"\nEpoch {epoch}/{EPOCHS_PHASE1}")
        print((
            f"[Model1] Train Loss: {train_loss1:.4f} | "
            f"Constraint Acc: {train_c_acc1:.4f}, Type Acc: {train_t_acc1:.4f} || "
            f"Val Loss: {val_loss1:.4f} | "
            f"Val Constraint Acc: {val_c_acc1:.4f}, Val Type Acc: {val_t_acc1:.4f}"
        ))
        print((
            f"[Model2] Train Loss: {train_loss2:.4f} | "
            f"Constraint Acc: {train_c_acc2:.4f}, Type Acc: {train_t_acc2:.4f} || "
            f"Val Loss: {val_loss2:.4f} | "
            f"Val Constraint Acc: {val_c_acc2:.4f}, Val Type Acc: {val_t_acc2:.4f}"
        ))

    ###################################
    # Phase 2: Pseudo-label unlabeled data (constraints only)
    ###################################
    print("\nPhase 2: Pseudo-labeling unlabeled data")
    df_unlabelled = pd.read_csv(UNLABELLED_FILE)
    unlabeled_texts = df_unlabelled["Text"].tolist()

    unlabeled_dataset = ConstraintDataset(
        unlabeled_texts, constraints=None, types=None, tokenizer=tokenizer, max_length=MAX_LENGTH
    )
    unlabeled_loader = DataLoader(unlabeled_dataset, BATCH_SIZE, pin_memory=pin_memory)

    preds1 = predict_constraint(model1, unlabeled_loader, device, USE_AMP)
    preds2 = predict_constraint(model2, unlabeled_loader, device, USE_AMP)

    pseudo_texts = []
    pseudo_constraints = []
    for txt, p1, p2 in zip(unlabeled_texts, preds1, preds2):
        if p1 == p2:
            pseudo_texts.append(txt)
            pseudo_constraints.append(p1)

    print(
        f"Pseudo-labeled {len(pseudo_texts)}/{len(unlabeled_texts)} "
        "where both models agree on constraint."
    )

    ###################################
    # Phase 3: Train a third classifier on combined data
    #   - Real labeled data: we have (constraint, type)
    #   - Pseudo-labeled data: we have (constraint, no type)
    ###################################
    print("\nPhase 3: Training third classifier with partial supervision")
    combined_texts = X_train + pseudo_texts
    combined_constraints = c_train + pseudo_constraints
    # For pseudo-labeled samples, we do not know type => store -1 => skip type loss
    combined_types = t_train + [-1]*len(pseudo_texts)

    # Minimal partial-supervision dataset
    class PartialDataset(Dataset):
        def __init__(self, texts, constraints, types, tokenizer, max_length=128):
            self.texts = texts
            self.constraints = constraints
            self.types = types
            self.tokenizer = tokenizer
            self.max_length = max_length

        def __len__(self):
            return len(self.texts)

        def __getitem__(self, idx):
            text = str(self.texts[idx])
            inputs = self.tokenizer(
                text,
                truncation=True,
                padding='max_length',
                max_length=self.max_length,
                return_tensors="pt"
            )
            input_ids = inputs["input_ids"].squeeze(0)
            attention_mask = inputs["attention_mask"].squeeze(0)
            c_label = torch.tensor(self.constraints[idx], dtype=torch.long)
            t_label = torch.tensor(self.types[idx], dtype=torch.long)
            return input_ids, attention_mask, c_label, t_label

    combined_dataset = PartialDataset(
        combined_texts, combined_constraints, combined_types, tokenizer, MAX_LENGTH
    )
    combined_loader = DataLoader(combined_dataset, BATCH_SIZE, shuffle=True, pin_memory=pin_memory)

    # Third classifier
    model3 = RobertaMultiTaskClassifier(NUM_CONSTRAINT_CLASSES, NUM_TYPE_CLASSES).to(device)
    optimizer3 = optim.AdamW(model3.parameters(), lr=LEARNING_RATE)

    for epoch in range(1, EPOCHS_PHASE3 + 1):
        train_loss, train_c_acc, train_t_acc = train_epoch_partial(
            model3, combined_loader, optimizer3, device, criterion, USE_AMP
        )
        # Evaluate on the val set (full labels)
        val_loss, val_c_acc, val_t_acc = eval_model(
            model3, val_loader, device, criterion, USE_AMP
        )
        print(f"\nEpoch {epoch}/{EPOCHS_PHASE3}")
        print((
            f"Model3 -> Train Loss: {train_loss:.4f} | "
            f"Constraint Acc: {train_c_acc:.4f}, Type Acc: {train_t_acc:.4f} || "
            f"Val Loss: {val_loss:.4f} | Val Constraint Acc: {val_c_acc:.4f}, Val Type Acc: {val_t_acc:.4f}"
        ))

    # Save final model
    torch.save(model3.state_dict(), MODEL_SAVE_PATH)
    print(f"\nModel weights saved to {MODEL_SAVE_PATH}")

if __name__ == "__main__":
    main()