SleepMaMi / train_utils.py

Keondopark

Checkpoint files converted to safetensors, README update

737ad15 2 days ago

15.9 kB

from configs.global_cfg import DATASET_PATH
import os
import pickle
import pandas as pd
from torch.optim.lr_scheduler import _LRScheduler
import math
import torch
from safetensors.torch import save_file
from lifelines.utils import concordance_index
from torchmetrics import Metric, MetricCollection
from torchmetrics.utilities import dim_zero_cat


def get_splits_from_file(ds, all_configs, data_root, split=None):
    sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects
    if split is None:
        sample_dir = os.path.join(data_root, DATASET_PATH[ds], sample_modality)
    else:
        sample_dir = os.path.join(data_root, DATASET_PATH[ds], split, sample_modality)

    chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
    info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
    with open(info_path, 'rb') as f:
        sig_info = pickle.load(f)

    subjects = sorted([sub for sub in sig_info['sig_info'].keys()])
    return subjects

def get_datasets(datasets, all_configs, data_root, splits=['train']):
    splits_set = dict()

    for split in splits:
        splits_set[split] = []

    for ds in datasets:
        # Use Physionet 2018 dataset        
        if ds == 'shhs1':
            def get_shhs_lists(shhs_split_file):
                ids = pd.read_csv(os.path.join("data","shhs_meta_data", shhs_split_file), header=None, names=['a'])
                ids = ids['a'].values.tolist()    
                list_files = ["shhs1-20"+id[11:15] for id in ids]
                return list_files
            list_files_val = get_shhs_lists("eval_list.txt")
            list_files_test = get_shhs_lists("test_list.txt")
            
            if 'train' in splits:
                train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
                missing_cases = ['shhs1-204709','shhs1-204806']

                train_subjects = set(train_subjects) - set(list_files_val) - set(list_files_test) - set(missing_cases)
                train_subjects = sorted(list(train_subjects)) 
                splits_set['train'].extend([('shhs1',f) for f in train_subjects])
            
            if 'val' in splits:
                val_subjects = get_splits_from_file(ds, all_configs, data_root, split='val')
                val_subjects = set(list_files_val) & set(val_subjects)
                val_subjects = sorted(list(val_subjects))
                splits_set['val'].extend([('shhs1',f) for f in val_subjects])
            
            if 'test' in splits:
                test_subjects = get_splits_from_file(ds, all_configs, data_root, split='test')
                test_subjects = set(list_files_test) & set(test_subjects)            
                test_subjects = sorted(list(test_subjects))
                splits_set['test'].extend([('shhs1',f) for f in test_subjects])
            
        elif ds == 'kiss' or 'cfs' in ds or 'isruc' in ds or 'sof' in ds:
            
            ds_main = 'kiss' if ds in 'kiss' else 'cfs' if 'cfs' in ds else 'isruc' if 'isruc' in ds else 'sof'            
            
            
            if 'train' in splits:
                train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
                splits_set['train'].extend([(ds_main,f) for f in train_subjects])
            
            if 'val' in splits:
                val_subjects = get_splits_from_file(ds, all_configs, data_root, split='val')
                splits_set['val'].extend([(ds_main,f) for f in val_subjects])
            
            if 'test' in splits:
                test_subjects = get_splits_from_file(ds, all_configs, data_root, split='test')
                splits_set['test'].extend([(ds_main,f) for f in test_subjects])
        elif ds in ['phy-tr', 'phy-te', 'mros1-val', 'mros1-te','mros2-val','mros2-te','mesa-val','mesa-te']:
            # sample_modality = all_configs.get('modalities', ['EEG'])[0]
            # sample_dir = os.path.join(data_root, DATASET_PATH[ds], sample_modality)
            # chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
            # info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
            # with open(info_path, 'rb') as f:
            #     sig_info = pickle.load(f)

            # all_subjects = sorted([sub for sub in sig_info['sig_info'].keys()])
            if 'train' in splits:
                train_subjects = get_splits_from_file(ds, all_configs, data_root, split=None)
                splits_set['train'].extend([(ds, f) for f in train_subjects])

        elif ds in ['shhs2', 'kiss','kvss','mros1','mros2','mesa']:
            if 'train' in splits:
                train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
                splits_set['train'].extend([(ds, f) for f in train_subjects])
            

            # splits_set['train'].extend([(ds, f) for f in all_subjects])

    return splits_set



def export_inference_weights(ckpt_path, model, out_path, metadata=None):
    """
    Strip a Lightning checkpoint down to the weights `model` loads and save them
    as safetensors (no optimizer / scheduler / callback state, no pickle).
    `model` is the inference backbone, e.g. `inference.build_micro_model()` or
    `inference.build_full_night_model()`. Keys are saved without the `model.`
    prefix; `pos_embed` is skipped since it is a fixed sin-cos table rebuilt at
    model construction. Keys absent from the checkpoint (e.g. `head` for a
    pretraining checkpoint) are not saved.
        from inference import build_micro_model
        export_inference_weights("experiments/micro.ckpt", build_micro_model(),
                                 "checkpoints/micro.safetensors", {"task": "micro"})
    """
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    src = ckpt["state_dict"] if "state_dict" in ckpt else ckpt

    weights = {}
    for k in model.state_dict():
        if "pos_embed" in k:
            continue
        if f"model.{k}" in src:
            weights[k] = src[f"model.{k}"].detach().clone().contiguous()
        elif k in src:
            weights[k] = src[k].detach().clone().contiguous()

    meta = {"source_epoch": str(ckpt.get("epoch")), "source_global_step": str(ckpt.get("global_step"))}
    meta.update({k: str(v) for k, v in (metadata or {}).items()})
    save_file(weights, out_path, metadata=meta)
    return weights


def get_layer_decay_param_groups(model, base_lr, weight_decay, layer_decay):
    """
    Create parameter groups with layer-wise learning rate decay.
    Args:
        model: The model whose parameters are grouped.
        base_lr: Base learning rate.
        weight_decay: Weight decay for regularization.
        layer_decay: Layer decay factor (e.g., 0.75).
    Returns:
        A list of parameter groups for the optimizer.
    """
    param_groups = []
    layers = list(model.children())
    num_layers = len(layers)

    # We use named_parameters to filter by name (e.g., A_log, dt, bias)
    # But we need to map them back to their layer index for LR decay.
    for i, layer in enumerate(layers):
        layer_scale = layer_decay ** (num_layers - i - 1)
        curr_lr = base_lr * layer_scale
        
        # We split parameters in each layer into 'decay' and 'no_decay'
        params_to_decay = []
        params_no_decay = []

        for name, param in layer.named_parameters():
            if not param.requires_grad:
                continue
            
            # Check if it's a Mamba-specific parameter or a standard "no-decay" candidate
            # Common Mamba names: A_log, D, dt_proj, x_proj, in_proj, conv1d
            # Standard no-decay: bias, norm weights
            is_ssm_param = any(ssm_name in name for ssm_name in ['A_log', 'D', 'dt', 'B', 'C'])
            is_bias_or_norm = name.endswith("bias") or "norm" in name
            
            if is_ssm_param or is_bias_or_norm:
                params_no_decay.append(param)
            else:
                params_to_decay.append(param)

        # Add groups for this specific layer
        if params_to_decay:
            param_groups.append({
                "params": params_to_decay,
                "lr": curr_lr,
                "weight_decay": weight_decay,
            })
        if params_no_decay:
            param_groups.append({
                "params": params_no_decay,
                "lr": curr_lr,
                "weight_decay": 0.0,
            })

    return param_groups

class WarmupCosineAnnealingLR(_LRScheduler):
    def __init__(self, optimizer, warmup_epochs, max_epochs, eta_min=1e-8, last_epoch=-1):
        """
        Step-based Warmup Cosine Annealing Learning Rate Scheduler.
        Args:
            optimizer: The optimizer for which to schedule the learning rate.
            warmup_epochs: Number of steps for the warmup phase.
            max_epochs: Total number of steps for the scheduler.
            eta_min: Minimum learning rate at the end of the schedule.
            last_epoch: The index of the last step (default: -1).
        """
        self.warmup_epochs = warmup_epochs
        self.max_epochs = max_epochs
        self.eta_min = eta_min
        super(WarmupCosineAnnealingLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        if self.last_epoch < self.warmup_epochs:
            # Linear warmup
            return [base_lr * (self.last_epoch + 1) / self.warmup_epochs for base_lr in self.base_lrs]
        else:
            # Cosine annealing
            progress = (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)
            return [self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * progress)) / 2 for base_lr in self.base_lrs]



def cox_ph_loss(hazards, event_times, is_event):
    # Sort event times and get corresponding indices for sorting other tensors
    event_times, sorted_idx = event_times.sort(dim=0, descending=True)
   

    # Make gathering robust: if sorted_idx is 1D and hazards/is_event are 2D,
    # use advanced indexing; otherwise fall back to gather which requires
    # the index tensor to have the same number of dimensions as input.
    if sorted_idx.dim() == 1 and hazards.dim() == 2:
        hazards = hazards[sorted_idx]
        is_event = is_event[sorted_idx]
    else:
        hazards = hazards.gather(0, sorted_idx)
        is_event = is_event.gather(0, sorted_idx)

    log_cumulative_hazard = torch.logcumsumexp(hazards.float(), dim=0)

    # Calculate losses for all labels simultaneously
    losses = (hazards - log_cumulative_hazard) * is_event
    losses = -losses  # Negative for maximization

    # Average loss per label
    label_loss = losses.sum(dim=0) / (is_event.sum(dim=0) + 1e-9)  # Avoid division by zero

    # Average across labels
    total_loss = label_loss.mean()

    return total_loss


def c_index_torch(scores, event_times, is_event):    
    return float(concordance_index(event_times, -scores, is_event))



class CIndex(Metric):
    def __init__(self, num_classes=6, **kwargs):
        super().__init__(**kwargs)
        self.add_state("scores", default=[], dist_reduce_fx="cat")
        self.add_state("event_time", default=[], dist_reduce_fx="cat")
        self.add_state("is_event", default=[], dist_reduce_fx="cat")
        self.num_classes = num_classes

    def update(self, score, t, e) -> None:
        self.scores.append(score)
        self.event_time.append(t)
        self.is_event.append(e)

    def compute(self):

        # parse inputs
        # Convert to float32 before numpy() since numpy doesn't support BFloat16
        scores = dim_zero_cat(self.scores).detach().cpu().float().numpy()
        event_time = dim_zero_cat(self.event_time).detach().cpu().float().numpy()
        is_event = dim_zero_cat(self.is_event).detach().cpu().float().numpy()
        
        from lifelines.utils import concordance_index
        cindex = []
        for i in range(self.num_classes):
            try:
                c = concordance_index(event_time[:,i], -scores[:,i], is_event[:,i])
                cindex.append(c)
            except (ZeroDivisionError, ValueError):
                # Not enough admissible pairs (common during sanity check or with small batches)
                # Return 0.5 (random predictor) as default
                cindex.append(0.5)
        return torch.tensor(cindex)
    
class AUROC(Metric):
    def __init__(self, num_classes=6, **kwargs):
        super().__init__(**kwargs)
        self.add_state("scores", default=[], dist_reduce_fx="cat")
        self.add_state("event_time", default=[], dist_reduce_fx="cat")
        self.add_state("is_event", default=[], dist_reduce_fx="cat")
        self.num_classes = num_classes
        self.measure_date = 365 * 6 # at 6 years

    def update(self, score, t, e) -> None:
        self.scores.append(score)
        self.event_time.append(t)
        self.is_event.append(e)

    def compute(self):

        # parse inputs
        # Convert to float32 before numpy() since numpy doesn't support BFloat16
        scores = dim_zero_cat(self.scores).detach().cpu().float().numpy()
        event_time = dim_zero_cat(self.event_time).detach().cpu().float().numpy()
        is_event = dim_zero_cat(self.is_event).detach().cpu().float().numpy()
        
        from sklearn.metrics import roc_auc_score, roc_curve
        # auroc = []
        # # For each disease, calculate AUROC
        # for i in range(self.num_classes):
        #     y_true = (event_time[:,i] <= self.measure_date).astype(int)
        #     y_scores = scores[:,i]
        #     try:
        #         c = roc_auc_score(y_true, y_scores)
        #         auroc.append(c)
        #     except (ZeroDivisionError, ValueError):
        #         # Not enough admissible pairs (common during sanity check or with small batches)
        #         # Return 0.5 (random predictor) as default
        #         auroc.append(0.5)
        # return torch.tensor(auroc)
        auroc = []
        # For each disease, calculate AUROC at 6-year follow-up (time-dependent)
        for i in range(self.num_classes):
            # Filter out zero times (invalid entries)
            valid_mask = event_time[:, i] > 0
            if valid_mask.sum() < 2:
                auroc.append(0.5)
                continue
            
            valid_scores = scores[valid_mask, i]
            valid_times = event_time[valid_mask, i]
            valid_events = is_event[valid_mask, i]
            
            # Time-dependent AUROC at measure_date:
            # - Samples with event before measure_date: y=1
            # - Samples with no event and observed until after measure_date: y=0
            # - Samples with censoring before measure_date: excluded (potentially informative censoring)
            
            # Consider only samples observed until measure_date
            observed_mask = valid_times >= self.measure_date
            # Or event occurred before measure_date
            event_before_mask = valid_events == 1
            
            # Use all samples with valid time info for AUROC
            # y_true: 1 if event before measure_date, 0 otherwise
            y_true = (valid_events == 1) & (valid_times <= self.measure_date)
            y_true = y_true.astype(int)
            y_scores = valid_scores
            
            try:
                # Check if there's variation in labels
                if len(np.unique(y_true)) < 2:
                    auroc.append(0.5)
                else:
                    c = roc_auc_score(y_true, y_scores)
                    auroc.append(c)
            except (ZeroDivisionError, ValueError) as e:
                # Not enough admissible pairs
                auroc.append(0.5)
        return torch.tensor(auroc)

	from configs.global_cfg import DATASET_PATH
	import os
	import pickle
	import pandas as pd
	from torch.optim.lr_scheduler import _LRScheduler
	import math
	import torch
	from safetensors.torch import save_file
	from lifelines.utils import concordance_index
	from torchmetrics import Metric, MetricCollection
	from torchmetrics.utilities import dim_zero_cat


	def get_splits_from_file(ds, all_configs, data_root, split=None):
	sample_modality = all_configs.get('modalities', ['EEG'])[0] # Make sure all modalities have the same number of subjects
	if split is None:
	sample_dir = os.path.join(data_root, DATASET_PATH[ds], sample_modality)
	else:
	sample_dir = os.path.join(data_root, DATASET_PATH[ds], split, sample_modality)

	chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
	info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
	with open(info_path, 'rb') as f:
	sig_info = pickle.load(f)

	subjects = sorted([sub for sub in sig_info['sig_info'].keys()])
	return subjects

	def get_datasets(datasets, all_configs, data_root, splits=['train']):
	splits_set = dict()

	for split in splits:
	splits_set[split] = []

	for ds in datasets:
	# Use Physionet 2018 dataset
	if ds == 'shhs1':
	def get_shhs_lists(shhs_split_file):
	ids = pd.read_csv(os.path.join("data","shhs_meta_data", shhs_split_file), header=None, names=['a'])
	ids = ids['a'].values.tolist()
	list_files = ["shhs1-20"+id[11:15] for id in ids]
	return list_files
	list_files_val = get_shhs_lists("eval_list.txt")
	list_files_test = get_shhs_lists("test_list.txt")

	if 'train' in splits:
	train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
	missing_cases = ['shhs1-204709','shhs1-204806']

	train_subjects = set(train_subjects) - set(list_files_val) - set(list_files_test) - set(missing_cases)
	train_subjects = sorted(list(train_subjects))
	splits_set['train'].extend([('shhs1',f) for f in train_subjects])

	if 'val' in splits:
	val_subjects = get_splits_from_file(ds, all_configs, data_root, split='val')
	val_subjects = set(list_files_val) & set(val_subjects)
	val_subjects = sorted(list(val_subjects))
	splits_set['val'].extend([('shhs1',f) for f in val_subjects])

	if 'test' in splits:
	test_subjects = get_splits_from_file(ds, all_configs, data_root, split='test')
	test_subjects = set(list_files_test) & set(test_subjects)
	test_subjects = sorted(list(test_subjects))
	splits_set['test'].extend([('shhs1',f) for f in test_subjects])

	elif ds == 'kiss' or 'cfs' in ds or 'isruc' in ds or 'sof' in ds:

	ds_main = 'kiss' if ds in 'kiss' else 'cfs' if 'cfs' in ds else 'isruc' if 'isruc' in ds else 'sof'


	if 'train' in splits:
	train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
	splits_set['train'].extend([(ds_main,f) for f in train_subjects])

	if 'val' in splits:
	val_subjects = get_splits_from_file(ds, all_configs, data_root, split='val')
	splits_set['val'].extend([(ds_main,f) for f in val_subjects])

	if 'test' in splits:
	test_subjects = get_splits_from_file(ds, all_configs, data_root, split='test')
	splits_set['test'].extend([(ds_main,f) for f in test_subjects])
	elif ds in ['phy-tr', 'phy-te', 'mros1-val', 'mros1-te','mros2-val','mros2-te','mesa-val','mesa-te']:
	# sample_modality = all_configs.get('modalities', ['EEG'])[0]
	# sample_dir = os.path.join(data_root, DATASET_PATH[ds], sample_modality)
	# chns = [f.split('.')[0] for f in os.listdir(sample_dir) if f.endswith('mmap')]
	# info_path = os.path.join(sample_dir, f'{chns[0]}.pickle')
	# with open(info_path, 'rb') as f:
	# sig_info = pickle.load(f)

	# all_subjects = sorted([sub for sub in sig_info['sig_info'].keys()])
	if 'train' in splits:
	train_subjects = get_splits_from_file(ds, all_configs, data_root, split=None)
	splits_set['train'].extend([(ds, f) for f in train_subjects])

	elif ds in ['shhs2', 'kiss','kvss','mros1','mros2','mesa']:
	if 'train' in splits:
	train_subjects = get_splits_from_file(ds, all_configs, data_root, split='train')
	splits_set['train'].extend([(ds, f) for f in train_subjects])


	# splits_set['train'].extend([(ds, f) for f in all_subjects])

	return splits_set



	def export_inference_weights(ckpt_path, model, out_path, metadata=None):
	"""
	Strip a Lightning checkpoint down to the weights `model` loads and save them
	as safetensors (no optimizer / scheduler / callback state, no pickle).

	`model` is the inference backbone, e.g. `inference.build_micro_model()` or
	`inference.build_full_night_model()`. Keys are saved without the `model.`
	prefix; `pos_embed` is skipped since it is a fixed sin-cos table rebuilt at
	model construction. Keys absent from the checkpoint (e.g. `head` for a
	pretraining checkpoint) are not saved.

	from inference import build_micro_model
	export_inference_weights("experiments/micro.ckpt", build_micro_model(),
	"checkpoints/micro.safetensors", {"task": "micro"})
	"""
	ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
	src = ckpt["state_dict"] if "state_dict" in ckpt else ckpt

	weights = {}
	for k in model.state_dict():
	if "pos_embed" in k:
	continue
	if f"model.{k}" in src:
	weights[k] = src[f"model.{k}"].detach().clone().contiguous()
	elif k in src:
	weights[k] = src[k].detach().clone().contiguous()

	meta = {"source_epoch": str(ckpt.get("epoch")), "source_global_step": str(ckpt.get("global_step"))}
	meta.update({k: str(v) for k, v in (metadata or {}).items()})
	save_file(weights, out_path, metadata=meta)
	return weights


	def get_layer_decay_param_groups(model, base_lr, weight_decay, layer_decay):
	"""
	Create parameter groups with layer-wise learning rate decay.
	Args:
	model: The model whose parameters are grouped.
	base_lr: Base learning rate.
	weight_decay: Weight decay for regularization.
	layer_decay: Layer decay factor (e.g., 0.75).
	Returns:
	A list of parameter groups for the optimizer.
	"""
	param_groups = []
	layers = list(model.children())
	num_layers = len(layers)

	# We use named_parameters to filter by name (e.g., A_log, dt, bias)
	# But we need to map them back to their layer index for LR decay.
	for i, layer in enumerate(layers):
	layer_scale = layer_decay ** (num_layers - i - 1)
	curr_lr = base_lr * layer_scale

	# We split parameters in each layer into 'decay' and 'no_decay'
	params_to_decay = []
	params_no_decay = []

	for name, param in layer.named_parameters():
	if not param.requires_grad:
	continue

	# Check if it's a Mamba-specific parameter or a standard "no-decay" candidate
	# Common Mamba names: A_log, D, dt_proj, x_proj, in_proj, conv1d
	# Standard no-decay: bias, norm weights
	is_ssm_param = any(ssm_name in name for ssm_name in ['A_log', 'D', 'dt', 'B', 'C'])
	is_bias_or_norm = name.endswith("bias") or "norm" in name

	if is_ssm_param or is_bias_or_norm:
	params_no_decay.append(param)
	else:
	params_to_decay.append(param)

	# Add groups for this specific layer
	if params_to_decay:
	param_groups.append({
	"params": params_to_decay,
	"lr": curr_lr,
	"weight_decay": weight_decay,
	})
	if params_no_decay:
	param_groups.append({
	"params": params_no_decay,
	"lr": curr_lr,
	"weight_decay": 0.0,
	})

	return param_groups

	class WarmupCosineAnnealingLR(_LRScheduler):
	def __init__(self, optimizer, warmup_epochs, max_epochs, eta_min=1e-8, last_epoch=-1):
	"""
	Step-based Warmup Cosine Annealing Learning Rate Scheduler.

	Args:
	optimizer: The optimizer for which to schedule the learning rate.
	warmup_epochs: Number of steps for the warmup phase.
	max_epochs: Total number of steps for the scheduler.
	eta_min: Minimum learning rate at the end of the schedule.
	last_epoch: The index of the last step (default: -1).
	"""
	self.warmup_epochs = warmup_epochs
	self.max_epochs = max_epochs
	self.eta_min = eta_min
	super(WarmupCosineAnnealingLR, self).__init__(optimizer, last_epoch)

	def get_lr(self):
	if self.last_epoch < self.warmup_epochs:
	# Linear warmup
	return [base_lr * (self.last_epoch + 1) / self.warmup_epochs for base_lr in self.base_lrs]
	else:
	# Cosine annealing
	progress = (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)
	return [self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * progress)) / 2 for base_lr in self.base_lrs]



	def cox_ph_loss(hazards, event_times, is_event):
	# Sort event times and get corresponding indices for sorting other tensors
	event_times, sorted_idx = event_times.sort(dim=0, descending=True)


	# Make gathering robust: if sorted_idx is 1D and hazards/is_event are 2D,
	# use advanced indexing; otherwise fall back to gather which requires
	# the index tensor to have the same number of dimensions as input.
	if sorted_idx.dim() == 1 and hazards.dim() == 2:
	hazards = hazards[sorted_idx]
	is_event = is_event[sorted_idx]
	else:
	hazards = hazards.gather(0, sorted_idx)
	is_event = is_event.gather(0, sorted_idx)

	log_cumulative_hazard = torch.logcumsumexp(hazards.float(), dim=0)

	# Calculate losses for all labels simultaneously
	losses = (hazards - log_cumulative_hazard) * is_event
	losses = -losses # Negative for maximization

	# Average loss per label
	label_loss = losses.sum(dim=0) / (is_event.sum(dim=0) + 1e-9) # Avoid division by zero

	# Average across labels
	total_loss = label_loss.mean()

	return total_loss


	def c_index_torch(scores, event_times, is_event):
	return float(concordance_index(event_times, -scores, is_event))



	class CIndex(Metric):
	def __init__(self, num_classes=6, **kwargs):
	super().__init__(**kwargs)
	self.add_state("scores", default=[], dist_reduce_fx="cat")
	self.add_state("event_time", default=[], dist_reduce_fx="cat")
	self.add_state("is_event", default=[], dist_reduce_fx="cat")
	self.num_classes = num_classes

	def update(self, score, t, e) -> None:
	self.scores.append(score)
	self.event_time.append(t)
	self.is_event.append(e)

	def compute(self):

	# parse inputs
	# Convert to float32 before numpy() since numpy doesn't support BFloat16
	scores = dim_zero_cat(self.scores).detach().cpu().float().numpy()
	event_time = dim_zero_cat(self.event_time).detach().cpu().float().numpy()
	is_event = dim_zero_cat(self.is_event).detach().cpu().float().numpy()

	from lifelines.utils import concordance_index
	cindex = []
	for i in range(self.num_classes):
	try:
	c = concordance_index(event_time[:,i], -scores[:,i], is_event[:,i])
	cindex.append(c)
	except (ZeroDivisionError, ValueError):
	# Not enough admissible pairs (common during sanity check or with small batches)
	# Return 0.5 (random predictor) as default
	cindex.append(0.5)
	return torch.tensor(cindex)

	class AUROC(Metric):
	def __init__(self, num_classes=6, **kwargs):
	super().__init__(**kwargs)
	self.add_state("scores", default=[], dist_reduce_fx="cat")
	self.add_state("event_time", default=[], dist_reduce_fx="cat")
	self.add_state("is_event", default=[], dist_reduce_fx="cat")
	self.num_classes = num_classes
	self.measure_date = 365 * 6 # at 6 years

	def update(self, score, t, e) -> None:
	self.scores.append(score)
	self.event_time.append(t)
	self.is_event.append(e)

	def compute(self):

	# parse inputs
	# Convert to float32 before numpy() since numpy doesn't support BFloat16
	scores = dim_zero_cat(self.scores).detach().cpu().float().numpy()
	event_time = dim_zero_cat(self.event_time).detach().cpu().float().numpy()
	is_event = dim_zero_cat(self.is_event).detach().cpu().float().numpy()

	from sklearn.metrics import roc_auc_score, roc_curve
	# auroc = []
	# # For each disease, calculate AUROC
	# for i in range(self.num_classes):
	# y_true = (event_time[:,i] <= self.measure_date).astype(int)
	# y_scores = scores[:,i]
	# try:
	# c = roc_auc_score(y_true, y_scores)
	# auroc.append(c)
	# except (ZeroDivisionError, ValueError):
	# # Not enough admissible pairs (common during sanity check or with small batches)
	# # Return 0.5 (random predictor) as default
	# auroc.append(0.5)
	# return torch.tensor(auroc)
	auroc = []
	# For each disease, calculate AUROC at 6-year follow-up (time-dependent)
	for i in range(self.num_classes):
	# Filter out zero times (invalid entries)
	valid_mask = event_time[:, i] > 0
	if valid_mask.sum() < 2:
	auroc.append(0.5)
	continue

	valid_scores = scores[valid_mask, i]
	valid_times = event_time[valid_mask, i]
	valid_events = is_event[valid_mask, i]

	# Time-dependent AUROC at measure_date:
	# - Samples with event before measure_date: y=1
	# - Samples with no event and observed until after measure_date: y=0
	# - Samples with censoring before measure_date: excluded (potentially informative censoring)

	# Consider only samples observed until measure_date
	observed_mask = valid_times >= self.measure_date
	# Or event occurred before measure_date
	event_before_mask = valid_events == 1

	# Use all samples with valid time info for AUROC
	# y_true: 1 if event before measure_date, 0 otherwise
	y_true = (valid_events == 1) & (valid_times <= self.measure_date)
	y_true = y_true.astype(int)
	y_scores = valid_scores

	try:
	# Check if there's variation in labels
	if len(np.unique(y_true)) < 2:
	auroc.append(0.5)
	else:
	c = roc_auc_score(y_true, y_scores)
	auroc.append(c)
	except (ZeroDivisionError, ValueError) as e:
	# Not enough admissible pairs
	auroc.append(0.5)
	return torch.tensor(auroc)