run_classification_task.py

import os
import json
import pandas as pd
import numpy as np

from dataloader.dataloader import SessionDataset, collate_fn, SessionDataset_KWA
from sklearn.metrics import confusion_matrix, multilabel_confusion_matrix

from model.event_embedding_model import EventEmbeddingModel
from model.session_embedding_model import SessionEmbeddingModel, SimpleSessionClassifier
from model.sequence_classifier import TransformerSessionClassifier
from model.loss import FocalLoss

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

import utils

from tqdm import tqdm

import os

def train_model(models, dataloader, optimizers, criterion, device):
    models['emb'].train()
    models['cls'].train()
    total_loss = 0

    all_logits = []
    all_labels = []
    total_loss = 0.0

    for batch in tqdm(dataloader, desc="Training"):
        optimizers['emb'].zero_grad()
        optimizers['cls'].zero_grad()

        x, y, mask = batch
        y = y.to(device)
        mask = mask.to(device)

        emb = models['emb'](x)
        logits = models['cls'](emb, mask)
        loss = criterion(logits, y)

        loss.backward()
        optimizers['emb'].step()
        optimizers['cls'].step()

        total_loss += loss.item()

        all_logits.append(logits.detach().cpu())
        all_labels.append(y.detach().cpu())

    best_thresholds, best_f1s = utils.find_best_threshold_per_class(
        torch.cat(all_logits,dim=0),
        torch.cat(all_labels,dim=0), metric='f1')

    return total_loss / len(dataloader), best_thresholds, best_f1s

def evaluate_model(models, dataloader, best_thresholds, device):
    models['emb'].eval()
    models['cls'].eval()
    best_thresholds = torch.tensor(best_thresholds, device=device).view(1, -1)  # shape: (1, C)
    best_thresholds = best_thresholds.squeeze(dim=1)

    all_preds, all_labels = [], []
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="Evaluating"):
            x, y, mask = batch
            y = y.to(device)
            mask = mask.to(device)
            emb = models['emb'](x)
            logits = models['cls'](emb, mask)
            preds = (torch.sigmoid(logits) > best_thresholds).long()
            all_preds.append(preds.cpu().numpy())
            all_labels.append(y.cpu().long().numpy())
    return np.concatenate(all_labels,axis=0), np.concatenate(all_preds,axis=0)


def main(config,target_flag: str, embedding_unit: str):
    print(target_flag)

    torch.manual_seed(2025)

    with open('train_test_split.json') as f:
        sp = json.load(f)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    num_epochs = config['num_epochs']

    for split_id in sp.keys():
        expdir = f'exp/{target_flag}/{embedding_unit}/{split_id}'
        os.makedirs(expdir,exist_ok=True)

        log_path = os.path.join(expdir, 'train.log')
        logger = utils.setup_logger(log_path=log_path)

        logger.info(target_flag)
        logger.info(split_id)

        if target_flag == 'in_context':
            train_dataset = SessionDataset(sp[split_id]['train'])
            test_dataset = SessionDataset(sp[split_id]['test'])
        elif target_flag == 'KWA':
            train_dataset = SessionDataset_KWA(sp[split_id]['train'],filter_incontext=True)
            test_dataset = SessionDataset_KWA(sp[split_id]['test'],filter_incontext=True)

        models = {}
        logger.info(f'embedding unit: {embedding_unit}')
        embedding_dim = config['text_embedding_dim']+config['cause_embed_dim']+config['app_embed_dim']
        if embedding_unit == 'session':
            models['emb'] = SessionEmbeddingModel(
                text_embedding_dim=config['text_embedding_dim'],
                cause_embed_dim=config['cause_embed_dim'],app_embed_dim=config['app_embed_dim'],
                output_dim=embedding_dim).to(device)
            models['cls'] = SimpleSessionClassifier(input_dim=embedding_dim, output_dim=train_dataset.get_dim()).to(device)
            bs_train = 16
        elif embedding_unit == 'event':
            models['emb'] = EventEmbeddingModel(
                text_embedding_dim=config['text_embedding_dim'],
                cause_embed_dim=config['cause_embed_dim'],app_embed_dim=config['app_embed_dim'],
                output_dim=embedding_dim).to(device)
            models['cls'] = TransformerSessionClassifier(input_dim=embedding_dim, output_dim=train_dataset.get_dim()).to(device)
            bs_train = 6
        optimizers = {}
        optimizers['emb'] = torch.optim.Adam(models['emb'].parameters(), lr=config['lr_emb'])
        optimizers['cls'] = torch.optim.Adam(models['cls'].parameters(), lr=config['lr_cls'])

        if config['loss_type'] == 'BCE':
            criterion = nn.BCEWithLogitsLoss()
        elif config['loss_type'] == 'BCEw':
            criterion = nn.BCEWithLogitsLoss(pos_weight=train_dataset.get_pos_weight().to(device))
        elif config['loss_type'] == 'Focal':
            criterion = FocalLoss(alpha=1.0, gamma=0.5)

        schedulers = {
            'emb': torch.optim.lr_scheduler.CosineAnnealingLR(optimizers['emb'], T_max=num_epochs),
            'cls': torch.optim.lr_scheduler.CosineAnnealingLR(optimizers['cls'], T_max=num_epochs),
        }

        train_loader = DataLoader(train_dataset, batch_size=bs_train, shuffle=True, collate_fn=collate_fn)
        test_loader = DataLoader(test_dataset, batch_size=16, collate_fn=collate_fn)

        for epoch in range(1,num_epochs+1):
            logger.info(f"Epoch {epoch}")

            train_loss, best_thres, best_f1s = train_model(models, train_loader, optimizers, criterion, device)

            for i, (t, f1) in enumerate(zip(best_thres, best_f1s)):
                logger.info(f"Class {i}: Best Threshold = {t:.2f}, Best F1 = {f1:.4f}")

            schedulers['emb'].step()
            schedulers['cls'].step()
            
            if not config['auto_thres']:
                best_thres = 0.5
            all_labels, all_preds = evaluate_model(models, test_loader, best_thresholds=best_thres, device=device)

            if target_flag == 'in_context':
                m = utils.eval_score(all_labels,all_preds)
                logger.info(f"Train Loss: {train_loss:.4f}, Test Acc: {m['acc']:.4f} Recall: {m['rec']:.4f} Precision {m['prec']:.4f} F1: {m['f1']:.4f}")

                cm = confusion_matrix(all_labels, all_preds)
                utils.plot_confusion_matrix(cm, expdir, epoch)
            else:
                active_classes = np.where(all_labels.sum(axis=0) > 0)[0]
                m = utils.eval_score(all_labels,all_preds,active_classes=active_classes)
                logger.info(f"Train Loss: {train_loss:.4f}, Test Jaccard_score: {m['jaccard_score']:.4f} Recall: {m['rec']:.4f} Precision {m['prec']:.4f} F1: {m['f1']:.4f}")

                cm = multilabel_confusion_matrix(all_labels, all_preds)
                
                utils.plot_multilabel_confusion_matrices(cm, train_dataset.get_labels(), expdir, epoch)            

            df = pd.DataFrame(all_labels, columns=[f"{name}_true" for name in train_dataset.get_labels()])
            df_preds = pd.DataFrame(all_preds, columns=[f"{name}_pred" for name in train_dataset.get_labels()])
            df_all = pd.concat([df, df_preds], axis=1)

            df_all.to_csv(f"{expdir}/predictions_epoch_{epoch}.csv", index=False)


if __name__ == '__main__':

    for target_flag in ['in_context', 'KWA']:
        for embedding_unit in ['session','event']:
            config = utils.load_config(f'config/config_{target_flag}.yaml')
            main(config,target_flag, embedding_unit=embedding_unit)