from .paths import BACKBONE import torch import torch.nn as nn import torch.nn.functional as F import math from transformers import AutoTokenizer from torch.optim import AdamW import numpy as np from sklearn.metrics import f1_score from typing import Tuple import lightning.pytorch as pl import gc from torch.optim.lr_scheduler import _LRScheduler from sklearn.metrics import f1_score, roc_auc_score, accuracy_score, matthews_corrcoef, precision_score, recall_score class GELU(nn.Module): def forward(self, x): return 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0))) class FocalLoss(nn.Module): """Focal loss. Sources: https://github.com/fcakyon/balanced-loss/blob/main/balanced_loss/losses.py https://github.com/vandit15/Class-balanced-loss-pytorch/blob/master/class_balanced_loss.py """ def __init__(self, gamma=2): super(FocalLoss, self).__init__() self.gamma = gamma def forward(self, inputs, targets, sample_weights): BCLoss = F.binary_cross_entropy_with_logits(input = inputs, target = targets,reduction = "none") if self.gamma == 0.0: modulator = 1.0 else: modulator = torch.exp(-self.gamma * targets * inputs - self.gamma * torch.log(1 + torch.exp(-1.0 * inputs))) loss = modulator * BCLoss weighted_loss = sample_weights * loss focal_loss = torch.sum(weighted_loss) focal_loss /= torch.sum(targets) return focal_loss class F1Loss(nn.Module): def __init__(self, epsilon=1e-7): super().__init__() self.epsilon = epsilon def forward(self, y_pred, y_true): assert y_pred.ndim == 2 assert y_true.ndim == 2 y_true = y_true.float() y_pred = torch.sigmoid(y_pred) tp = (y_true * y_pred).sum(dim=0) fp = ((1 - y_true) * y_pred).sum(dim=0) fn = (y_true * (1 - y_pred)).sum(dim=0) precision = tp / (tp + fp + self.epsilon) recall = tp / (tp + fn + self.epsilon) f1 = 2 * (precision * recall) / (precision + recall + self.epsilon) f1 = f1.clamp(min=self.epsilon, max=1-self.epsilon) return 1 - f1.mean() def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(x, cos, sin): # x: (B, L, H, D); cos and sin: (1, L, D). cos = cos.unsqueeze(2) sin = sin.unsqueeze(2) return (x * cos) + (rotate_half(x) * sin) class RotaryEmbedding(torch.nn.Module): """Rotary position embeddings (Su et al.), inspired by GPT-NeoX. https://arxiv.org/abs/2104.09864 https://github.com/ZhuiyiTechnology/roformer https://github.com/EleutherAI/gpt-neox """ def __init__(self, dim: int, *_, **__): super().__init__() inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq) self._seq_len_cached = None self._cos_cached = None self._sin_cached = None def _update_cos_sin_tables(self, x, seq_dimension=1): seq_len = x.shape[seq_dimension] if seq_len != self._seq_len_cached or self._cos_cached.device != x.device: self._seq_len_cached = seq_len t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq) freqs = torch.einsum("i,j->ij", t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1).to(x.device) self._cos_cached = emb.cos()[None, :, :] self._sin_cached = emb.sin()[None, :, :] return self._cos_cached, self._sin_cached def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k) return ( apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached), ) def macro_f1(y_true, y_pred, thresholds): y_pred_binary = (y_pred >= thresholds).float() tp = (y_true * y_pred_binary).sum(dim=0) fp = ((1 - y_true) * y_pred_binary).sum(dim=0) fn = (y_true * (1 - y_pred_binary)).sum(dim=0) precision = tp / (tp + fp + 1e-7) recall = tp / (tp + fn + 1e-7) f1 = 2 * precision * recall / (precision + recall + 1e-7) macro_f1 = f1.mean() return macro_f1 def safeguard_softmax(logits, dim=-1): max_logits, _ = logits.max(dim=dim, keepdim=True) exp_logits = torch.exp(logits - max_logits) exp_sum = exp_logits.sum(dim=dim, keepdim=True) probs = exp_logits / (exp_sum + 1e-7) return probs class PositionalAttentionHead(nn.Module): def __init__(self, hidden_dim, n_heads): super(PositionalAttentionHead, self).__init__() self.n_heads = n_heads self.hidden_dim = hidden_dim self.head_dim = hidden_dim // n_heads self.preattn_ln = nn.LayerNorm(self.head_dim) self.Q = nn.Linear(self.head_dim, self.head_dim, bias=False) self.K = nn.Linear(self.head_dim, self.head_dim, bias=False) self.V = nn.Linear(self.head_dim, self.head_dim, bias=False) self.rot_emb = RotaryEmbedding(self.head_dim) def forward(self, x, attention_mask): batch_size, seq_len, _ = x.size() x = x.view(batch_size, seq_len, self.n_heads, self.head_dim) x = self.preattn_ln(x) q = self.Q(x) k = self.K(x) v = self.V(x) q, k = self.rot_emb(q, k) gc.collect() torch.cuda.empty_cache() attn_scores = torch.einsum('bqhd,bkhd->bhqk', q, k) / math.sqrt(self.head_dim) attn_scores = attn_scores.masked_fill(torch.logical_not(attention_mask.unsqueeze(1).unsqueeze(1)), float("-inf")) attn_probs = safeguard_softmax(attn_scores, dim=-1) x = torch.einsum('bhqk,bkhd->bqhd', attn_probs, v) x = x.reshape(batch_size, seq_len, self.hidden_dim) gc.collect() torch.cuda.empty_cache() return x, attn_probs class CosineAnnealingWithWarmup(_LRScheduler): def __init__(self, optimizer, warmup_steps, total_steps, eta_ratio=0.1, last_epoch=-1): self.warmup_steps = warmup_steps self.total_steps = total_steps self.eta_ratio = eta_ratio super(CosineAnnealingWithWarmup, self).__init__(optimizer, last_epoch) def get_lr(self): if self.last_epoch < self.warmup_steps: return [base_lr * self.last_epoch / self.warmup_steps for base_lr in self.base_lrs] progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps) cosine_decay = 0.5 * (1 + np.cos(np.pi * progress)) decayed_lr = (1 - self.eta_ratio) * cosine_decay + self.eta_ratio return [decayed_lr * base_lr for base_lr in self.base_lrs] class AdaptiveMarginTripletLoss(nn.Module): def __init__(self, initial_margin=0.05, scale_factor=0.1, min_margin=0.05, max_margin=2.0): super(AdaptiveMarginTripletLoss, self).__init__() self.initial_margin = initial_margin self.scale_factor = scale_factor self.min_margin = min_margin self.max_margin = max_margin def forward(self, anchor, positive, negative): pos_dist = F.pairwise_distance(anchor, positive, p=2) neg_dist = F.pairwise_distance(anchor, negative, p=2) adaptive_margins = self.initial_margin + self.scale_factor * (neg_dist - pos_dist) adaptive_margins = torch.clamp(adaptive_margins, min=self.min_margin, max=self.max_margin) losses = F.relu(pos_dist - neg_dist + adaptive_margins) return losses.mean() class RobertaLMHead(nn.Module): def __init__(self, embed_dim, output_dim, weight): super().__init__() self.dense = nn.Linear(embed_dim, embed_dim) self.layer_norm = nn.LayerNorm(embed_dim) self.weight = weight self.gelu = GELU() self.bias = nn.Parameter(torch.zeros(output_dim)) def forward(self, features): x = self.dense(features) x = self.gelu(x) x = self.layer_norm(x) x = F.linear(x, self.weight) + self.bias return x def pad_to_max_length(tensor, max_length, pad_value=1): return F.pad(tensor, (0, max_length - tensor.shape[1]), value=pad_value) def compute_ece(probs, labels, n_bins=10): """Per-class calibration error for probabilities and labels of shape (B, C).""" bin_boundaries = torch.linspace(0, 1, n_bins + 1) bin_lowers = bin_boundaries[:-1] bin_uppers = bin_boundaries[1:] n_classes = probs.size(1) ece_per_class = torch.zeros(n_classes, device=probs.device) for class_idx in range(n_classes): class_probs = probs[:, class_idx] class_labels = labels[:, class_idx] ece = torch.zeros(1, device=probs.device) for bin_lower, bin_upper in zip(bin_lowers, bin_uppers): in_bin = (class_probs.gt(bin_lower.item()) * class_probs.le(bin_upper.item())).float() prop_in_bin = in_bin.mean() if prop_in_bin > 0: accuracy_in_bin = (class_labels * in_bin).sum() / in_bin.sum() avg_confidence_in_bin = (class_probs * in_bin).sum() / in_bin.sum() ece += torch.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin ece_per_class[class_idx] = ece.item() return ece_per_class class MultitaskProteinModel(pl.LightningModule): def __init__(self, hparams, esm_model, num_layers_to_finetune, input_dim=1280, output_dim=12, class_weights=None, hidden_dim=512, n_heads=10): super(MultitaskProteinModel, self).__init__() self.save_hyperparameters(hparams) self.esm_model = esm_model for param in self.esm_model.parameters(): param.requires_grad = False for i in range(num_layers_to_finetune): for param in self.esm_model.encoder.layer[-i-1].parameters(): param.requires_grad = True self.lm_head = RobertaLMHead(embed_dim = 1280, output_dim=33, weight=self.esm_model.embeddings.word_embeddings.weight) self.attn_head = PositionalAttentionHead(input_dim, n_heads) self.attn_ln = nn.LayerNorm(input_dim) self.attn_skip = nn.Linear(input_dim, input_dim) self.linear_layers = nn.ModuleList() for _ in range(self.hparams.num_linear_layers): self.linear_layers.append(nn.Linear(input_dim, input_dim)) self.reduction_layers = nn.Sequential( nn.Linear(input_dim, hidden_dim), GELU(), nn.Linear(hidden_dim, output_dim) ) self.clf_ln = nn.LayerNorm(input_dim) self.class_weights = class_weights self.classification_thresholds = nn.Parameter(torch.tensor([0.5]*output_dim)) self.historical_memory = 0.9 self.training_step_outputs = [] self.training_step_true_labels = [] self.testing_step_predictions = [] self.testing_step_labels = [] self.test_dataloader = None self.tokenizer = AutoTokenizer.from_pretrained(str(BACKBONE), local_files_only=True) self.corruption_epoch_threshold = 5 def setup(self, stage=None): self.__build_loss() def __build_loss(self): self.contrastive_loss = AdaptiveMarginTripletLoss() focal_gamma = torch.tensor(self.hparams.focal_gamma).to(self.device) self.class_weights = self.class_weights.to(self.device) self.focal_loss = FocalLoss(gamma=focal_gamma) self.f1_loss = F1Loss() self.reconstruction_loss = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id) def forward(self, input_ids, attention_masks=None): outputs = self.esm_model(input_ids=input_ids, attention_mask=attention_masks, output_hidden_states=True) embeddings = outputs.last_hidden_state attention_masks = attention_masks x_pool, x_attns = self.attn_head(embeddings, attention_masks) x_pool = self.attn_ln(x_pool + self.attn_skip(x_pool)) for linear_layer in self.linear_layers: residue = x_pool x_pool = linear_layer(x_pool) x_pool = F.silu(x_pool) x_pool = x_pool + residue x_weighted = torch.einsum('bhlk,bld->bhld', x_attns, x_pool) x_combined = x_weighted.mean(dim=1) x_combined = self.clf_ln(x_combined) mlm_logits = self.lm_head(x_combined) attention_masks = attention_masks.unsqueeze(-1).float() attention_sum = attention_masks.sum(dim=1, keepdim=True) x_combined_masked = (x_combined * attention_masks).sum(dim=1) / attention_sum.squeeze(1) x_pred = self.reduction_layers(x_combined_masked) gc.collect() torch.cuda.empty_cache() return x_pred, x_attns, x_combined_masked, mlm_logits def forward_from_embeds(self, inputs_embeds, attention_masks=None): esm = self.esm_model # Disable token dropout for embedding input. if getattr(esm.embeddings, "token_dropout", False): esm.embeddings.token_dropout = False outs = esm( inputs_embeds=inputs_embeds, attention_mask=attention_masks.to(torch.long), output_hidden_states=True, ) embeddings = outs.last_hidden_state x_pool, x_attns = self.attn_head(embeddings, attention_masks) x_pool = self.attn_ln(x_pool + self.attn_skip(x_pool)) for linear_layer in self.linear_layers: residue = x_pool x_pool = F.silu(linear_layer(x_pool)) x_pool = x_pool + residue x_weighted = torch.einsum('bhlk,bld->bhld', x_attns, x_pool).mean(dim=1) x_weighted = self.clf_ln(x_weighted) mask_f = attention_masks.unsqueeze(-1).float() x_combined_masked = (x_weighted * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp_min(1e-6) x_pred = self.reduction_layers(x_combined_masked) return x_pred, x_attns, x_combined_masked, None def predict(self, input_ids, attention_masks=None): x_pred, _, _, _ = self.forward(input_ids, attention_masks) classification_output = torch.sigmoid(x_pred) predictions = (classification_output >= self.classification_thresholds).float() for i, pred in enumerate(predictions): if pred.sum() == 0: weighted_probs = classification_output[i] max_class = torch.argmax(weighted_probs) predictions[i, max_class] = 1.0 confidence_scores = torch.abs(classification_output - self.classification_thresholds) / self.classification_thresholds max_confidence_scores = confidence_scores.max(dim=1).values return classification_output, predictions, max_confidence_scores def training_step(self, batch, batch_idx): anchor_input_ids, anchor_attention_masks = batch['anchor_input_ids'], batch['anchor_attention_mask'] positive_input_ids, positive_attention_masks = batch['positive_input_ids'], batch['positive_attention_mask'] negative_input_ids, negative_attention_masks = batch['negative_input_ids'], batch['negative_attention_mask'] anchor_labels = batch['anchor_labels'] positive_labels = batch['positive_labels'] negative_labels = batch['negative_labels'] anchor_x_pred, _, anchor_x_combined, mlm_logits = self.forward(anchor_input_ids, anchor_attention_masks) positive_x_pred, _, positive_x_combined, _ = self.forward(positive_input_ids, positive_attention_masks) negative_x_pred, _, negative_x_combined, _ = self.forward(negative_input_ids, negative_attention_masks) reconstruction_loss = self.reconstruction_loss(mlm_logits.view(-1, mlm_logits.size(-1)), anchor_input_ids.view(-1)) anchor_focal_loss = self.focal_loss(anchor_x_pred, anchor_labels, self.class_weights) positive_focal_loss = self.focal_loss(positive_x_pred, positive_labels, self.class_weights) semi_hard_negative_ce_loss = F.binary_cross_entropy_with_logits(input=negative_x_pred, target=negative_labels) focal_loss = (anchor_focal_loss + positive_focal_loss) / 2 anchor_f1_loss = self.f1_loss(anchor_x_pred, anchor_labels) positive_f1_loss = self.f1_loss(positive_x_pred, positive_labels) semi_hard_f1_loss = self.f1_loss(negative_x_pred, negative_labels) semi_hard_triplet_loss = self.contrastive_loss(anchor_x_combined, positive_x_combined, negative_x_combined) if self.current_epoch > self.corruption_epoch_threshold: hard_negative_input_ids, hard_negative_attention_masks = batch['hard_negative_input_ids'], batch['hard_negative_attention_mask'] hard_negative_labels = batch['hard_negative_labels'] if self.current_epoch < 15: epoch_ratio = min(1.0, (self.current_epoch - self.corruption_epoch_threshold) / (self.hparams.epochs - self.corruption_epoch_threshold)) hard_weight = epoch_ratio semi_hard_weight = 1 - hard_weight else: semi_hard_weight = 0.0 hard_weight = 1.0 hard_negative_x_pred, _, hard_negative_x_combined, _ = self.forward(hard_negative_input_ids, hard_negative_attention_masks) hard_negative_ce_loss = F.binary_cross_entropy_with_logits(input=hard_negative_x_pred, target=hard_negative_labels) hard_triplet_loss = self.contrastive_loss(anchor_x_combined, positive_x_combined, hard_negative_x_combined) hard_f1_loss = self.f1_loss(hard_negative_x_pred, hard_negative_labels) negative_ce_loss = semi_hard_weight * semi_hard_negative_ce_loss + hard_weight * hard_negative_ce_loss triplet_loss = semi_hard_weight * semi_hard_triplet_loss + hard_weight * hard_triplet_loss f1_loss = anchor_f1_loss + positive_f1_loss + semi_hard_weight * semi_hard_triplet_loss + hard_weight * hard_f1_loss else: negative_ce_loss = semi_hard_negative_ce_loss triplet_loss = semi_hard_triplet_loss f1_loss = anchor_f1_loss + positive_f1_loss + semi_hard_f1_loss loss = focal_loss * 2 + negative_ce_loss + f1_loss + triplet_loss * 2 + reconstruction_loss self.training_step_outputs.append(anchor_x_pred) self.training_step_outputs.append(positive_x_pred) self.training_step_outputs.append(negative_x_pred) self.training_step_true_labels.append(anchor_labels) self.training_step_true_labels.append(positive_labels) self.training_step_true_labels.append(negative_labels) macro_f1_val = macro_f1(anchor_labels, torch.sigmoid(anchor_x_pred), self.classification_thresholds) self.log('train_macro_f1', macro_f1_val, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) gc.collect() torch.cuda.empty_cache() gc.collect() torch.cuda.empty_cache() self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) return loss def on_train_epoch_end(self): epoch_outputs = torch.cat(self.training_step_outputs, dim=0) epoch_true_labels = torch.cat(self.training_step_true_labels, dim=0) epoch_outputs = torch.sigmoid(epoch_outputs) self.update_class_thresholds(epoch_outputs, epoch_true_labels) threshold_dict = {f"threshold_{i}": self.classification_thresholds[i].item() for i in range(15)} self.log_dict(threshold_dict, on_epoch=True) self.training_step_outputs.clear() self.training_step_true_labels.clear() gc.collect() torch.cuda.empty_cache() def validation_step(self, batch, batch_idx): anchor_input_ids, anchor_attention_masks = batch['anchor_input_ids'], batch['anchor_attention_mask'] positive_input_ids, positive_attention_masks = batch['positive_input_ids'], batch['positive_attention_mask'] negative_input_ids, negative_attention_masks = batch['negative_input_ids'], batch['negative_attention_mask'] anchor_labels = batch['anchor_labels'] positive_labels = batch['positive_labels'] negative_labels = batch['negative_labels'] anchor_x_pred, _, anchor_x_combined, mlm_logits = self.forward(anchor_input_ids, anchor_attention_masks) positive_x_pred, _, positive_x_combined, _ = self.forward(positive_input_ids, positive_attention_masks) negative_x_pred, _, negative_x_combined, _ = self.forward(negative_input_ids, negative_attention_masks) reconstruction_loss = self.reconstruction_loss(mlm_logits.view(-1, mlm_logits.size(-1)), anchor_input_ids.view(-1)) anchor_focal_loss = self.focal_loss(anchor_x_pred, anchor_labels , self.class_weights) positive_focal_loss = self.focal_loss(positive_x_pred, positive_labels, self.class_weights) negative_ce_loss = F.binary_cross_entropy_with_logits(input=negative_x_pred, target=negative_labels) focal_loss = (anchor_focal_loss + positive_focal_loss) / 2 f1_loss = self.f1_loss(anchor_x_pred, anchor_labels) triplet_loss = self.contrastive_loss(anchor_x_combined, positive_x_combined, negative_x_combined) loss = focal_loss * 2 + negative_ce_loss + f1_loss + triplet_loss * 2 + reconstruction_loss macro_f1_val = macro_f1(anchor_labels, torch.sigmoid(anchor_x_pred), self.classification_thresholds) self.log('val_macro_f1', macro_f1_val, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) gc.collect() torch.cuda.empty_cache() self.log('val_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) def test_step(self, batch, batch_idx): input_ids, labels, attention_masks = batch['input_ids'], batch['labels'], batch['attention_masks'] x_pred, _, _ = self.forward_test(input_ids, attention_masks) predictions = torch.sigmoid(x_pred) self.testing_step_predictions.append(predictions) self.testing_step_labels.append(labels) def on_test_epoch_end(self,): predictions = torch.cat(self.testing_step_predictions, dim=0) labels = torch.cat(self.testing_step_labels, dim=0) predictions_binary = (predictions >= self.classification_thresholds).float().cpu().numpy() labels_binary = labels.cpu().numpy() f1 = f1_score(labels_binary, predictions_binary, average='macro') roc_auc = roc_auc_score(labels_binary, predictions.cpu().to(torch.float32).numpy(), average='micro') accuracy = accuracy_score(labels_binary.flatten(), predictions_binary.flatten()) mcc = matthews_corrcoef(labels_binary.flatten(), predictions_binary.flatten()) precision = precision_score(labels_binary, predictions_binary, average='macro') recall = recall_score(labels_binary, predictions_binary, average='macro') results = { 'test_F1_macro': f1, 'test_AUCROC': roc_auc, 'test_Accuracy': accuracy, 'test_MCC': mcc, 'test_Precision': precision, 'test_Recall': recall, } self.log_dict(results) self.testing_step_predictions.clear() self.testing_step_labels.clear() gc.collect() torch.cuda.empty_cache() def configure_optimizers(self): optimizer = AdamW(self.parameters(), lr=self.hparams.lr, betas=(0.9, 0.95), weight_decay=1e-5) schedulers = CosineAnnealingWithWarmup(optimizer, warmup_steps=100, eta_ratio=0.1, total_steps=9000) lr_schedulers = { "scheduler": schedulers, "name": 'learning_rate_logs', "interval": 'step', 'frequency': 1 } return [optimizer], [lr_schedulers] def update_class_thresholds(self, inputs, targets): with torch.no_grad(): min_threshold_value = 0.001 thresholds = torch.linspace(0.0, 1.0, steps=1000, device=inputs.device) num_classes = targets.size(1) best_f1_scores = torch.zeros(num_classes, device=inputs.device) best_thresholds = torch.full((num_classes,), min_threshold_value, device=inputs.device) for threshold in thresholds: binary_predictions = (inputs >= threshold).float() tp = (binary_predictions * targets).sum(dim=0) fp = (binary_predictions * (1 - targets)).sum(dim=0) fn = ((1 - binary_predictions) * targets).sum(dim=0) precision = tp / (tp + fp + 1e-7) recall = tp / (tp + fn + 1e-7) f1_scores = 2 * precision * recall / (precision + recall + 1e-7) is_better = f1_scores > best_f1_scores best_f1_scores = torch.where(is_better, f1_scores, best_f1_scores) best_thresholds = torch.where(is_better, threshold.expand_as(best_f1_scores), best_thresholds) updated_thresholds = self.historical_memory * self.classification_thresholds.data + (1 - self.historical_memory) * best_thresholds self.classification_thresholds = nn.Parameter(torch.clamp(updated_thresholds, min=min_threshold_value)) gc.collect() torch.cuda.empty_cache() class ResidueEncoder(MultitaskProteinModel): """Return per-residue clf_ln outputs for diffusion training.""" def forward(self, input_ids, attention_masks=None): captured = [] hook = self.clf_ln.register_forward_hook(lambda module, inputs, output: captured.append(output)) try: logits, attention, pooled, mlm_logits = super().forward(input_ids, attention_masks) finally: hook.remove() return logits, attention, captured[0], mlm_logits