BioM3-PenCL push with no weights

Stage1_source/PL_wrapper.py

@@ -0,0 +1,1613 @@
+# pytorch fucntions
+import torch
+from torch import nn, optim
+from torch.nn import functional as F
+import torch.distributed as dist
+
+# PL functions
+import pytorch_lightning as pl
+from pytorch_lightning import Trainer, seed_everything
+
+# misc functions
+import itertools
+import matplotlib.pyplot as plt
+import numpy as np
+import sys
+from tqdm import tqdm
+import time
+
+# other learning packages
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
+
+# our packages
+import Stage1_source.helper_funcs as helper_tools
+import Stage1_source.preprocess as prep
+import Stage1_source.model as mod
+
+
+######################
+# Default PL wrapper #
+######################
+
+class PL_PEN_CL(pl.LightningModule):
+
+
+    def __init__(
+            self,
+            args: any,
+            model: nn.Module,
+            text_tokenizer: any,
+            sequence_tokenizer: any
+        ):
+
+        super().__init__()
+        # arguments
+        self.script_args = args
+        
+        # model components
+        self.model = model
+
+        # tokenizers
+        self.text_tokenizer = text_tokenizer
+        self.sequence_tokenizer = sequence_tokenizer
+        
+        # validation tracker for outputs
+        self.val_text_joint_latents = []
+        self.val_seq_joint_latents = []
+
+        # prediction tracker for outputs
+        self.predict_text_joint_latents = []
+        self.predict_seq_joint_latents = []
+            
+    def forward(
+            self,
+            x_t: torch.Tensor,
+            x_s: torch.Tensor
+        ) -> (
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor
+        ):
+
+        outputs = self.model(
+                        x_t=x_t,
+                        x_s=x_s
+        )
+
+        return (
+                outputs['text_joint_latent'],
+                outputs['seq_joint_latent'],
+        )
+
+    def training_step(
+            self,
+            batch: torch.Tensor,
+            batch_idx: any,
+        ) -> dict:
+         
+        if isinstance(batch, list):
+            # split the 
+            text_batch, protein_batch = batch
+        
+        # forward pass 
+        z_t, z_s = self(
+                    x_t=text_batch,
+                    x_s=protein_batch
+        )
+        dist.barrier()
+
+        # gather all tensors
+        z_t_all = self.all_gather(z_t, sync_grads=True)
+        dist.barrier()
+        z_s_all = self.all_gather(z_s, sync_grads=True)
+        
+        # stack the embeddings
+        z_t_all = z_t_all.view(-1, z_t.shape[-1])
+        z_s_all = z_s_all.view(-1, z_s.shape[-1])
+      
+        # compute loss values
+        loss, logits = self.model.compute_loss(
+                protein_embeddings=z_s_all,
+                text_embeddings=z_t_all
+        )
+        
+        # track loss ...
+        self.log('train_loss', loss, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        
+        # track metrics
+        metric_dict = self.performance_metrics(logits=logits)
+        for key in metric_dict:
+            values = metric_dict[key]
+
+            final_key = 'train_' + key
+            self.log(final_key, metric_dict[key], prog_bar=True if 'f1' in key else False, on_step=True, on_epoch=True, sync_dist=True)
+        
+        if batch_idx == 0:
+            gpu_memory_usage = helper_tools.print_gpu_initialization()
+            self.log(f'gpu_memory_usage', gpu_memory_usage, sync_dist=True)
+
+        return {'loss': loss}
+
+
+    def validation_step(
+            self,
+            batch: list,
+            batch_idx: any
+        ) -> dict:
+
+        # split the batch
+        if isinstance(batch, list):
+            # mean loss
+            text_batch, protein_batch = batch
+        
+        # forward pass 
+        z_t, z_s = self(
+                    x_t=text_batch,
+                    x_s=protein_batch
+        )
+        
+        dist.barrier()
+        # gather all tensors
+        z_t_all = self.all_gather(z_t, sync_grads=True).view(-1, z_t.shape[-1])
+        dist.barrier()
+        z_s_all = self.all_gather(z_s, sync_grads=True).view(-1, z_s.shape[-1])
+        
+        # stack the embeddings
+        z_t_all = z_t_all.view(-1, z_t.shape[-1])
+        z_s_all = z_s_all.view(-1, z_s.shape[-1])
+
+        # compute loss values
+        loss, logits = self.model.compute_loss(
+                protein_embeddings=z_s_all,
+                text_embeddings=z_t_all
+        )
+        
+        
+        # track validation loss ...
+        self.log('valid_loss', loss, prog_bar=True, sync_dist=True)
+
+        # copmute validation metrics
+        metric_dict = self.performance_metrics(logits=logits.detach().cpu())
+
+        for key in metric_dict:
+            values = metric_dict[key]
+            final_key = 'valid_' + key
+            self.log(final_key, metric_dict[key], prog_bar=True if 'f1' in key else False, sync_dist=True)
+        
+        # collect joint embedding
+        self.val_text_joint_latents.append(z_t_all.detach().cpu())
+        self.val_seq_joint_latents.append(z_s_all.detach().cpu())
+
+        return {'valid_loss': loss}
+
+    def on_validation_epoch_end(self):
+        
+        # collect and aggregate outputs from all validation steps
+        val_z_t_joint = torch.cat(self.val_text_joint_latents, dim=0)
+        val_z_s_joint = torch.cat(self.val_seq_joint_latents, dim=0)
+        
+        # compute singular values 
+        text_log_sigma_k, S_text = self.compute_singular(val_z_t_joint.detach().cpu())
+        protein_log_sigma_k, S_protein = self.compute_singular(val_z_s_joint.detach().cpu())
+        
+        # save image pngs for tracking dimensionality collapse
+        self.save_png_to_tensorboard(
+                data=text_log_sigma_k.numpy(),
+                title='text',
+        )
+        self.save_png_to_tensorboard(
+                data=protein_log_sigma_k.numpy(),
+                title='protein'
+        )
+
+        # free memory
+        self.val_text_joint_latents.clear()
+        self.val_seq_joint_latents.clear()
+        
+
+        # compute effective rank (RankME):
+        erank_text = self.compute_effective_rank(sigma_ks=S_text)
+        erank_protein = self.compute_effective_rank(sigma_ks=S_protein)
+        
+        # log erank metrics
+        self.log('valid_erank_text', erank_text, sync_dist=True)
+        self.log('valid_erank_protein', erank_protein, sync_dist=True)
+
+
+    def configure_optimizers(self,):
+
+        params = [
+                {"params": self.model.protein_encoder.parameters(), "lr": self.script_args.protein_encoder_lr},
+                {"params": self.model.text_encoder.parameters(), "lr": self.script_args.text_encoder_lr},
+                {"params": itertools.chain(
+                    self.model.protein_projection.parameters(),
+                    self.model.text_projection.parameters()
+                    ),
+                "lr": self.script_args.head_lr,
+                "weight_decay": self.script_args.weight_decay}
+        ]
+
+        optimizer = torch.optim.AdamW(params, weight_decay=self.script_args.weight_decay)
+
+        return {
+                "optimizer": optimizer,
+        }
+    
+    @torch.no_grad()
+    def compute_class_metrics(
+            self,
+            outputs: torch.Tensor,
+            targets: torch.Tensor,
+            source: str
+        ) -> dict:
+
+        # convert torch tensors to numpy array
+        outputs_np = outputs.numpy()
+        targets_np = targets.numpy()
+
+        # compute the metrics
+        accuracy = accuracy_score(targets_np, outputs_np.round())
+        precision = precision_score(targets_np, outputs_np.round(), average='micro')
+        recall = recall_score(targets_np, outputs_np.round(), average='micro')
+        f1 = f1_score(targets_np, outputs_np.round(), average='micro')
+
+        return {
+                f'{source}_accuracy': accuracy,
+                f'{source}_precision': precision,
+                f'{source}_recall': recall,
+                f'{source}_f1': f1
+        }
+
+    @torch.no_grad()
+    def performance_metrics(self, logits: torch.Tensor) -> tuple:
+
+        logits = logits.cpu().float()
+
+        # get probs
+        p_text = F.softmax(logits, dim=-1) # prob of a given text captions aligning well with seq. pairs
+        p_seq = F.softmax(logits.T, dim=-1) # prob of a given seq aligning well with text pairs
+        p_tot = (p_seq + p_text) / 2 # total prob
+
+        # get class labels
+        y_pred_text = torch.argmax(p_text, dim=-1)
+        y_pred_seq = torch.argmax(p_seq, dim=-1)
+        y_pred = torch.argmax(p_tot, dim=-1)
+        y_true = torch.arange(y_pred_text.shape[0])
+        
+        # compute class metrics
+        text_metrics = self.compute_class_metrics(
+                                    outputs=y_pred_text,
+                                    targets=y_true,
+                                    source='text'
+        )
+        seq_metrics = self.compute_class_metrics(
+                                    outputs=y_pred_seq,
+                                    targets=y_true,
+                                    source='seq'
+        )
+        total_metrics = self.compute_class_metrics(
+                                    outputs=y_pred,
+                                    targets=y_true,
+                                    source='total'
+        )
+
+        # combine dicts into one
+        combined_dict = {}
+        combined_dict.update(text_metrics)
+        combined_dict.update(seq_metrics)
+        combined_dict.update(total_metrics)
+
+        return combined_dict
+    
+    @torch.no_grad()
+    def compute_singular(self, inputs: torch.Tensor) -> (
+            torch.Tensor,
+            torch.Tensor
+        ):
+
+        # goal of this function: track for dimensionality collapse
+        # inputs dim: (batch_size, emb_dim)
+
+        mean_inputs = torch.mean(inputs, dim=0) # average over batch dimension
+        norm_inputs = inputs - mean_inputs # normalize vectors
+       
+       # compute correlation matrix  #TODO: double check work...
+        C = torch.zeros((norm_inputs.shape[-1], norm_inputs.shape[-1]))
+        for sample_idx in range(norm_inputs.shape[0]):
+            norm_vector = norm_inputs[sample_idx, :].unsqueeze(0)
+            C += norm_vector.T @ norm_vector
+        C *= 1/norm_vector.shape[0]
+
+        _, S, _ = torch.linalg.svd(C, full_matrices=False)
+        
+        # return singular value indexes 
+        log_sigma_k, _ = torch.sort(torch.log(S), descending=True)
+        return (
+                log_sigma_k,
+                S
+        )
+        
+    def compute_effective_rank(self, sigma_ks: torch.Tensor) -> torch.Tensor:
+        """
+        references:
+            - Roy et al. The effective rank: a measure of effective dimensionality
+            - Garrido et al. RankMe: Assessing the Downstream Performnace of Pretrained SS Reps by their Rank.
+        """
+        # sort the singular values
+        sigma_ks, _ = torch.sort(sigma_ks, descending=True)
+        
+        # copute L1 norm for sing values.
+        l1_norm_sigma = torch.norm(sigma_ks, p=1)
+        
+        # compute singular value distribution
+        p_k = sigma_ks / l1_norm_sigma  + torch.finfo(torch.float).eps
+        
+        # compute Shannon entropy
+        entropy = - torch.sum(p_k * torch.log(p_k))
+    
+        # get effective rank (RankME):
+        erank = torch.exp(entropy)
+
+        return erank
+
+    def save_png_to_tensorboard(
+            self,
+            data: np.single,
+            title: str,
+            x_axis_label: str='Singular Value Rank Index',
+            y_axis_label: str='Log of singular values',
+            ):
+    
+        current_epoch = self.trainer.current_epoch
+        
+        # Plot the line
+        fig, ax = plt.subplots(dpi=300)
+        ax.plot(data)
+        ax.set_xlabel(x_axis_label)
+        ax.set_ylabel(y_axis_label)
+        ax.set_title(title)
+        ax.set_ylim([-25,3])
+
+        # Log the plot in TensorBoard 
+        self.logger.experiment.add_figure(f'{title}_SingularValues_{current_epoch}', fig, current_epoch)
+
+    def predict_step(
+            self,
+            batch: torch.Tensor,
+            batch_idx: torch.Tensor,
+            dataloder_idx: bool=False
+        ) -> (  
+                torch.Tensor,
+                torch.Tensor
+        ):
+
+
+        if isinstance(batch, list):
+            # mean loss
+            text_batch, protein_batch = batch
+            outputs = self(
+                    x_t=text_batch,
+                    x_s=protein_batch,
+            )
+        
+        z_t_joint, z_p_joint = outputs
+
+        self.predict_text_joint_latents.append(z_t_joint.detach().cpu())
+        self.predict_seq_joint_latents.append(z_p_joint.detach().cpu())
+
+        return outputs
+
+    def on_predict_epoch_end(self, outputs=None):
+
+        self.predict_text_joint_latents = torch.cat(self.predict_text_joint_latents).cpu()
+        self.predict_seq_joint_latents = torch.cat(self.predict_seq_joint_latents).cpu()
+
+
+
+##########################
+# Masked-task PL wrapper #
+##########################
+
+class mask_PL_PEN_CL(pl.LightningModule):
+
+
+    def __init__(
+            self,
+            args: any,
+            model: nn.Module,
+            text_tokenizer: any,
+            sequence_tokenizer: any
+    ):
+
+        super().__init__()
+        # arguments
+        self.script_args = args
+        
+        # model components
+        self.model = model
+    
+        # tokenizers
+        self.text_tokenizer = text_tokenizer
+        self.sequence_tokenizer = sequence_tokenizer
+        
+        # validation tracker for outputs
+        self.val_text_joint_latents = []
+        self.val_seq_joint_latents = []
+    
+        # prediction tracker for outputs
+        self.predict_text_joint_latents = []
+        self.predict_seq_joint_latents = []
+
+    def forward(
+            self,
+            x_t: torch.Tensor,
+            x_s: torch.Tensor,
+            compute_masked_logits: bool=False
+        ) -> (
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor
+        ):
+
+        outputs = self.model(
+                        x_t=x_t,
+                        x_s=x_s,
+                        compute_masked_logits=compute_masked_logits
+        )
+        
+        if compute_masked_logits:
+            # forward pass for computing logits for masked language objective
+            return (
+                    outputs['text_masked_logits'],
+                    outputs['protein_masked_logits']
+            )
+        else:
+            # forward pass for computing latent embeddings in the joint space
+            return (
+                outputs['text_joint_latent'],
+                outputs['seq_joint_latent'],
+            )
+
+    def training_step(
+            self,
+            batch: torch.Tensor,
+            batch_idx: any,
+        ) -> dict:
+         
+        if isinstance(batch, list):
+            # split the data
+            text_batch, protein_batch, text_mask_batch, protein_mask_batch = batch
+        
+        # forward pass 
+        z_t, z_s = self(
+                    x_t=text_batch,
+                    x_s=protein_batch,
+                    compute_masked_logits=False
+        )
+        dist.barrier()
+
+        # gather all tensors
+        z_t_all = self.all_gather(z_t, sync_grads=True)
+        dist.barrier()
+        z_s_all = self.all_gather(z_s, sync_grads=True)
+        
+        # stack the embeddings
+        z_t_all = z_t_all.view(-1, z_t.shape[-1])
+        z_s_all = z_s_all.view(-1, z_s.shape[-1])
+      
+        # compute loss values
+        loss_align, logits = self.model.compute_loss(
+                protein_embeddings=z_s_all,
+                text_embeddings=z_t_all
+        )
+        
+        # compute mask language model logits
+        logits_t_mask, logits_s_mask = self(
+                x_t=text_mask_batch,
+                x_s=protein_mask_batch,
+                compute_masked_logits=True
+        )
+        
+        # compute mask language loss for biomedical expert model
+        loss_text_mask = self.model.compute_masked_lang_loss(
+                logits_masked=logits_t_mask,
+                targets=text_batch,
+                targets_masked=text_mask_batch,
+                mask_token_id=self.text_tokenizer.mask_token_id
+        )
+        
+        # compute mask language loss for protein expert model
+        loss_sequence_mask = self.model.compute_masked_lang_loss(
+                logits_masked=logits_s_mask,
+                targets=protein_batch,
+                targets_masked=protein_mask_batch,
+                mask_token_id=self.sequence_tokenizer.mask_idx
+        )
+        
+        
+        # total loss
+        loss = loss_align + loss_text_mask + loss_sequence_mask
+
+        # track loss ...
+        self.log('train_loss', loss, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('train_loss_align', loss_align, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('train_loss_text_mask', loss_text_mask, prog_bar=False, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('train_loss_seq_mask', loss_sequence_mask, prog_bar=False, on_step=True, on_epoch=True, sync_dist=True)
+
+        # track metrics
+        metric_dict = self.performance_metrics(logits=logits)
+        for key in metric_dict:
+            values = metric_dict[key]
+
+            final_key = 'train_' + key
+            self.log(final_key, metric_dict[key], prog_bar=True if 'f1' in key else False, on_step=True, on_epoch=True, sync_dist=True)
+        
+        if batch_idx == 0:
+            gpu_memory_usage = helper_tools.print_gpu_initialization()
+            self.log(f'gpu_memory_usage', gpu_memory_usage, sync_dist=True)
+
+        return {'loss': loss}
+
+
+    def validation_step(
+            self,
+            batch: list,
+            batch_idx: any
+        ) -> dict:
+
+        # split the batch
+        if isinstance(batch, list):
+            # mean loss
+            text_batch, protein_batch, text_mask_batch, protein_mask_batch = batch
+        
+        # forward pass 
+        z_t, z_s = self(
+                    x_t=text_batch,
+                    x_s=protein_batch
+        )
+        
+        dist.barrier()
+        # gather all tensors
+        z_t_all = self.all_gather(z_t, sync_grads=True).view(-1, z_t.shape[-1])
+        dist.barrier()
+        z_s_all = self.all_gather(z_s, sync_grads=True).view(-1, z_s.shape[-1])
+        
+        # stack the embeddings
+        z_t_all = z_t_all.view(-1, z_t.shape[-1])
+        z_s_all = z_s_all.view(-1, z_s.shape[-1])
+
+        # compute loss values
+        loss_align, logits = self.model.compute_loss(
+                protein_embeddings=z_s_all,
+                text_embeddings=z_t_all
+        )
+        
+        # compute mask language model logits
+        logits_t_mask, logits_s_mask = self(
+                x_t=text_mask_batch,
+                x_s=protein_mask_batch,
+                compute_masked_logits=True
+        )
+        
+        # compute mask language loss for biomedical expert model
+        loss_text_mask = self.model.compute_masked_lang_loss(
+                logits_masked=logits_t_mask,
+                targets=text_batch,
+                targets_masked=text_mask_batch,
+                mask_token_id=self.text_tokenizer.mask_token_id
+        )
+        
+        # compute mask language loss for protein expert model
+        loss_sequence_mask = self.model.compute_masked_lang_loss(
+                logits_masked=logits_s_mask,
+                targets=protein_batch,
+                targets_masked=protein_mask_batch,
+                mask_token_id=self.sequence_tokenizer.mask_idx
+        )
+        
+        # total loss
+        loss = loss_align + loss_text_mask + loss_sequence_mask
+
+        # track validation loss ...
+        self.log('valid_loss', loss, prog_bar=True, sync_dist=True)
+        self.log('valid_loss_align', loss_align, prog_bar=True, sync_dist=True)
+        self.log('valid_loss_text_mask', loss_text_mask, prog_bar=False, sync_dist=True)
+        self.log('valid_loss_seq_mask', loss_sequence_mask, prog_bar=False, sync_dist=True)
+
+        # copmute validation metrics
+        metric_dict = self.performance_metrics(logits=logits.detach().cpu())
+
+        for key in metric_dict:
+            values = metric_dict[key]
+            final_key = 'valid_' + key
+            self.log(final_key, metric_dict[key], prog_bar=True if 'f1' in key else False, sync_dist=True)
+        
+        # collect joint embedding
+        self.val_text_joint_latents.append(z_t_all.detach().cpu())
+        self.val_seq_joint_latents.append(z_s_all.detach().cpu())
+
+        return {'valid_loss': loss}
+
+    def on_validation_epoch_end(self):
+        
+ #       # collect and aggregate outputs from all validation steps
+ #      val_z_t_joint = torch.cat(self.val_text_joint_latents, dim=0)
+ #       val_z_s_joint = torch.cat(self.val_seq_joint_latents, dim=0)
+        
+        # compute singular values 
+ #       text_log_sigma_k, S_text = self.compute_singular(val_z_t_joint.detach().cpu())
+ #       protein_log_sigma_k, S_protein = self.compute_singular(val_z_s_joint.detach().cpu())
+        
+        # save image pngs for tracking dimensionality collapse
+ #       self.save_png_to_tensorboard(
+ #               data=text_log_sigma_k.numpy(),
+ #               title='text',
+ #       )
+ #       self.save_png_to_tensorboard(
+ #               data=protein_log_sigma_k.numpy(),
+ #               title='protein'
+ #       )
+
+        # free memory
+        self.val_text_joint_latents.clear()
+        self.val_seq_joint_latents.clear()
+        
+
+        # compute effective rank (RankME):
+ #       erank_text = self.compute_effective_rank(sigma_ks=S_text)
+ #       erank_protein = self.compute_effective_rank(sigma_ks=S_protein)
+        
+        # log erank metrics
+ #       self.log('valid_erank_text', erank_text, sync_dist=True)
+ #       self.log('valid_erank_protein', erank_protein, sync_dist=True)
+
+
+    def configure_optimizers(self,):
+
+        params = [
+                {"params": self.model.protein_encoder.parameters(), "lr": self.script_args.protein_encoder_lr},
+                {"params": self.model.text_encoder.parameters(), "lr": self.script_args.text_encoder_lr},
+                {"params": itertools.chain(
+                    self.model.protein_projection.parameters(),
+                    self.model.text_projection.parameters()
+                    ),
+                "lr": self.script_args.head_lr,
+                "weight_decay": self.script_args.weight_decay}
+        ]
+
+        optimizer = torch.optim.AdamW(params, weight_decay=self.script_args.weight_decay)
+
+        return {
+                "optimizer": optimizer,
+        }
+    
+    @torch.no_grad()
+    def compute_class_metrics(
+            self,
+            outputs: torch.Tensor,
+            targets: torch.Tensor,
+            source: str
+        ) -> dict:
+
+        # convert torch tensors to numpy array
+        outputs_np = outputs.numpy()
+        targets_np = targets.numpy()
+
+        # compute the metrics
+        accuracy = accuracy_score(targets_np, outputs_np.round())
+        precision = precision_score(targets_np, outputs_np.round(), average='micro')
+        recall = recall_score(targets_np, outputs_np.round(), average='micro')
+        f1 = f1_score(targets_np, outputs_np.round(), average='micro')
+
+        return {
+                f'{source}_accuracy': accuracy,
+                f'{source}_precision': precision,
+                f'{source}_recall': recall,
+                f'{source}_f1': f1
+        }
+
+    @torch.no_grad()
+    def performance_metrics(self, logits: torch.Tensor) -> tuple:
+
+        logits = logits.cpu().float()
+
+        # get probs
+        p_text = F.softmax(logits, dim=-1) # prob of a given text captions aligning well with seq. pairs
+        p_seq = F.softmax(logits.T, dim=-1) # prob of a given seq aligning well with text pairs
+        p_tot = (p_seq + p_text) / 2 # total prob
+
+        # get class labels
+        y_pred_text = torch.argmax(p_text, dim=-1)
+        y_pred_seq = torch.argmax(p_seq, dim=-1)
+        y_pred = torch.argmax(p_tot, dim=-1)
+        y_true = torch.arange(y_pred_text.shape[0])
+        
+        # compute class metrics
+        text_metrics = self.compute_class_metrics(
+                                    outputs=y_pred_text,
+                                    targets=y_true,
+                                    source='text'
+        )
+        seq_metrics = self.compute_class_metrics(
+                                    outputs=y_pred_seq,
+                                    targets=y_true,
+                                    source='seq'
+        )
+        total_metrics = self.compute_class_metrics(
+                                    outputs=y_pred,
+                                    targets=y_true,
+                                    source='total'
+        )
+
+        # combine dicts into one
+        combined_dict = {}
+        combined_dict.update(text_metrics)
+        combined_dict.update(seq_metrics)
+        combined_dict.update(total_metrics)
+
+        return combined_dict
+    
+    @torch.no_grad()
+    def compute_singular(self, inputs: torch.Tensor) -> (
+            torch.Tensor,
+            torch.Tensor
+        ):
+
+        # goal of this function: track for dimensionality collapse
+        # inputs dim: (batch_size, emb_dim)
+
+        mean_inputs = torch.mean(inputs, dim=0) # average over batch dimension
+        norm_inputs = inputs - mean_inputs # normalize vectors
+       
+       # compute correlation matrix  #TODO: double check work...
+        C = torch.zeros((norm_inputs.shape[-1], norm_inputs.shape[-1]))
+        for sample_idx in tqdm(range(norm_inputs.shape[0])):
+            norm_vector = norm_inputs[sample_idx, :].unsqueeze(0)
+            C += norm_vector.T @ norm_vector
+        C *= 1/norm_vector.shape[0]
+
+        _, S, _ = torch.linalg.svd(C, full_matrices=False)
+        
+        # return singular value indexes 
+        log_sigma_k, _ = torch.sort(torch.log(S), descending=True)
+        return (
+                log_sigma_k,
+                S
+        )
+        
+    def compute_effective_rank(self, sigma_ks: torch.Tensor) -> torch.Tensor:
+        """
+        references:
+            - Roy et al. The effective rank: a measure of effective dimensionality
+            - Garrido et al. RankMe: Assessing the Downstream Performnace of Pretrained SS Reps by their Rank.
+        """
+        # sort the singular values
+        sigma_ks, _ = torch.sort(sigma_ks, descending=True)
+        
+        # copute L1 norm for sing values.
+        l1_norm_sigma = torch.norm(sigma_ks, p=1)
+        
+        # compute singular value distribution
+        p_k = sigma_ks / l1_norm_sigma  + torch.finfo(torch.float).eps
+        
+        # compute Shannon entropy
+        entropy = - torch.sum(p_k * torch.log(p_k))
+    
+        # get effective rank (RankME):
+        erank = torch.exp(entropy)
+
+        return erank
+
+    def save_png_to_tensorboard(
+            self,
+            data: np.single,
+            title: str,
+            x_axis_label: str='Singular Value Rank Index',
+            y_axis_label: str='Log of singular values',
+            ):
+    
+        current_epoch = self.trainer.current_epoch
+        
+        # Plot the line
+        fig, ax = plt.subplots(dpi=300)
+        ax.plot(data)
+        ax.set_xlabel(x_axis_label)
+        ax.set_ylabel(y_axis_label)
+        ax.set_title(title)
+        ax.set_ylim([-25,3])
+
+        # Log the plot in TensorBoard 
+        self.logger.experiment.add_figure(f'{title}_SingularValues_{current_epoch}', fig, current_epoch)
+
+    def predict_step(
+            self,
+            batch: torch.Tensor,
+            batch_idx: torch.Tensor,
+            dataloder_idx: bool=False
+        ) -> (  
+                torch.Tensor,
+                torch.Tensor
+        ):
+
+
+        if isinstance(batch, list):
+            # mean loss
+            text_batch, protein_batch = batch
+            outputs = self(
+                    x_t=text_batch,
+                    x_s=protein_batch,
+                    compute_masked_logits=False
+            )
+        
+        z_t_joint, z_p_joint = outputs
+
+        self.predict_text_joint_latents.append(z_t_joint.detach().cpu())
+        self.predict_seq_joint_latents.append(z_p_joint.detach().cpu())
+
+        return outputs
+
+    def on_predict_epoch_end(self, outputs=None):
+
+        self.predict_text_joint_latents = torch.cat(self.predict_text_joint_latents).cpu()
+        self.predict_seq_joint_latents = torch.cat(self.predict_seq_joint_latents).cpu()
+
+
+
+########################
+# Pfam-task PL wrapper #
+########################
+
+
+class pfam_PL_PEN_CL(pl.LightningModule):
+
+
+    def __init__(
+            self,
+            args: any,
+            model: nn.Module,
+            text_tokenizer: any,
+            sequence_tokenizer: any
+    ):
+
+        super().__init__()
+        # arguments
+        self.script_args = args
+        
+        # model components
+        self.model = model
+    
+        # tokenizers
+        self.text_tokenizer = text_tokenizer
+        self.sequence_tokenizer = sequence_tokenizer
+        
+        # validation tracker for outputs
+        self.val_text_joint_latents = []
+        self.val_seq_joint_latents = []
+    
+        # predictions...
+        self.predict_text_joint_latents = [] 
+        self.predict_seq_joint_latents = []
+
+    def forward(
+            self,
+            x_t: torch.Tensor,
+            x_p: torch.Tensor,
+            compute_masked_logits: bool=False
+        ) -> (
+                torch.Tensor,
+                torch.Tensor,
+                torch.Tensor
+        ):
+
+        outputs = self.model(
+                        x_t=x_t,
+                        x_s=x_p,
+                        compute_masked_logits=compute_masked_logits
+        )
+        
+        if compute_masked_logits:
+            # forward pass for computing logits for masked language objective
+            return (
+                    outputs['text_masked_logits'],
+                    outputs['protein_masked_logits']
+            )
+        else:
+            # forward pass for computing latent embeddings in the joint space
+            return (
+                outputs['text_joint_latent'],
+                outputs['seq_joint_latent'],
+            )
+
+
+    def on_train_batch_start(self, batch, batch_idx):
+        self.batch_start_time = time.time()
+
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        batch_end_time = time.time()
+        batch_time = batch_end_time - self.batch_start_time
+        #print(f'Rank={dist.get_rank()}: time to process batch is {batch_time}')
+        #self.log(f'batch_time_rank_{dist.get_rank()}', batch_time, on_step=True, on_epoch=False)
+
+    def training_step(self, batch: torch.Tensor, batch_idx: any) -> dict:
+        """
+        Execute a single training step.
+
+        Given a batch of data, this function processes both Swiss-Prot and Pfam data through the model, computes 
+        various loss values including inter-modal, intra-modal, and masked language model losses for both text 
+        and protein sequences. This function also computes and logs various metrics and GPU memory usage.
+
+        Parameters:
+        - batch: The input data batch. This can include multiple types of data.
+        - batch_idx: Index of the current batch.
+
+        Steps:
+        1. Split the data into Swiss-Prot and Pfam batches, if the batch is a list.
+        2. Forward pass the Swiss-Prot data through the model.
+        3. Synchronize and gather embeddings from all GPUs.
+        4. Forward pass the Pfam data through the model.
+        5. Synchronize and gather Pfam embeddings from all GPUs.
+        6. Concatenate Swiss-Prot and Pfam embeddings.
+        7. Compute inter-modal and intra-modal loss values.
+        8. Compute masked language model logits for the concatenated batch.
+        9. Compute masked language loss for both text and protein sequences.
+        10. Compute and log the total loss and individual loss components.
+        11. Compute and log performance metrics.
+        12. Log GPU memory usage at the start of training.
+
+        Returns:
+        - Dictionary containing the total loss value.
+
+        Note: 
+        This function is intended to be used within a distributed (multi-GPU) training context, as evident 
+        from the use of barriers and gathering operations. It's designed to handle batches that contain both 
+        Swiss-Prot and Pfam data, both being biological datasets used in multi-modal protein embeddings. 
+        The function utilizes both inter-modal (between modalities) and intra-modal (within the same modality) 
+        contrastive losses, as well as masked language modeling objectives similar to BERT's MLM objective.
+        """
+
+        # Check if the batch is a list and split data if so.
+        if isinstance(batch, list):
+            text_batch, protein_batch, text_mask_batch, protein_mask_batch, \
+            pfam_text_batch, pfam_protein_batch, pfam_text_mask_batch, pfam_protein_mask_batch, \
+            bool_pfam_vector = batch
+    
+
+        #print(f'rank={dist.get_rank()}: text size {text_batch.shape}')
+        
+        #start_time_forward_pass = time.time()
+        # Forward pass with Swiss-Prot data.
+        z_t_swiss, z_p_swiss = self(
+            x_t=text_batch,
+            x_p=protein_batch,
+            compute_masked_logits=False
+        )
+        # Timer end and log
+        #end_time_forward_pass = time.time()
+        #print(f"Rank={dist.get_rank()}: Time taken for Swiss-Prot forward pass: {end_time_forward_pass - start_time_forward_pass} seconds.")
+
+        # Ensure all GPUs are synchronized.
+        dist.barrier()
+
+        # Forward pass with Pfam data.
+        z_t_pfam, z_p_pfam = self(
+            x_t=pfam_text_batch,
+            x_p=pfam_protein_batch,
+            compute_masked_logits=False
+        )
+        dist.barrier()
+        
+        #Gather tensors from all GPUs.
+        z_t_swiss_all = self.all_gather(z_t_swiss, sync_grads=True)
+        dist.barrier()
+        z_p_swiss_all = self.all_gather(z_p_swiss, sync_grads=True)
+
+        # Reshape the embeddings.
+        z_t_swiss_all = z_t_swiss_all.view(-1, z_t_swiss.shape[-1])
+        z_p_swiss_all = z_p_swiss_all.view(-1, z_p_swiss.shape[-1])
+
+
+        # Gather tensors from all GPUs.
+        z_t_pfam_all = self.all_gather(z_t_pfam, sync_grads=True)
+        dist.barrier()
+        z_p_pfam_all = self.all_gather(z_p_pfam, sync_grads=True)
+
+        # Reshape the embeddings.
+        z_t_pfam_all = z_t_pfam_all.view(-1, z_t_pfam.shape[-1])
+        z_p_pfam_all = z_p_pfam_all.view(-1, z_p_pfam.shape[-1])
+
+        # Concatenate Swiss-Prot and Pfam embeddings.
+        z_t_all = torch.cat((z_t_swiss_all, z_t_pfam_all), dim=0)
+        z_p_all = torch.cat((z_p_swiss_all, z_p_pfam_all), dim=0)
+        
+        # Timer start
+        #start_time_loss_computation = time.time()
+
+        # Compute inter-modal loss.
+        loss_align, logits = self.model.compute_inter_loss(
+            protein_embeddings=z_p_all,
+            text_embeddings=z_t_all,
+            batch_size=z_p_all.shape[0] // 2
+        )
+        # Timer end and log
+        #end_time_loss_computation = time.time()
+        #print(f"Rank={dist.get_rank()}: Time taken for loss computation: {end_time_loss_computation - start_time_loss_computation} seconds.")
+
+
+        # Compute intra-modal loss.
+        loss_intra, cosine_similarity = self.model.compute_intra_loss(
+            protein_embeddings=z_p_all,
+            batch_size=z_p_all.shape[0] // 2
+        )
+
+        # Concatenate batches for masked language modeling.
+        all_text_batch = torch.cat((text_batch, pfam_text_batch), dim=0)
+        all_protein_batch = torch.cat((protein_batch, pfam_protein_batch), dim=0)
+        all_text_mask_batch = torch.cat((text_mask_batch, pfam_text_mask_batch), dim=0)
+        all_protein_mask_batch = torch.cat((protein_mask_batch, pfam_protein_mask_batch), dim=0)
+        
+        #TODO: timer start
+        #start_time_mask_comp = time.time()
+
+        # Compute masked language model logits.
+        logits_t_mask, logits_s_mask = self(
+            x_t=all_text_mask_batch,
+            x_p=all_protein_mask_batch,
+            compute_masked_logits=True
+        )
+        #end_time_mask_comp = time.time()
+        #print(f"Rank={dist.get_rank()}: Time taken for mask predictions: {end_time_mask_comp - start_time_mask_comp} seconds.")
+
+
+        # Compute masked language model loss for text data.
+        loss_text_mask = self.model.compute_masked_lang_loss(
+            logits_masked=logits_t_mask,
+            targets=all_text_batch,
+            targets_masked=all_text_mask_batch,
+            mask_token_id=self.text_tokenizer.mask_token_id
+        )
+
+        # Compute masked language model loss for protein data.
+        loss_sequence_mask = self.model.compute_masked_lang_loss(
+            logits_masked=logits_s_mask,
+            targets=all_protein_batch,
+            targets_masked=all_protein_mask_batch,
+            mask_token_id=self.sequence_tokenizer.mask_idx
+        )
+
+
+        if self.script_args.dataset_type == 'pfam':
+            # Aggregate all computed losses.
+            loss = loss_align + loss_intra + loss_text_mask + loss_sequence_mask
+        
+        elif self.script_args.dataset_type == 'pfam_ablated':
+            # Aggregate all losses besides PFC.
+            loss = loss_align + loss_text_mask + loss_sequence_mask
+        else:
+            # Add an assertion here
+            assert self.script_args.dataset_type in ['pfam', 'pfam_ablated'], "Unexpected dataset_type value"
+            sys.stderr.write("Unexpected dataset_type value\n")
+            sys.exit(1)
+
+        # Log the individual and total loss values.
+        self.log('train_loss', loss, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('train_loss_align', loss_align, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('train_loss_intra', loss_intra, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('train_loss_text_mask', loss_text_mask, prog_bar=False, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('train_loss_seq_mask', loss_sequence_mask, prog_bar=False, on_step=True, on_epoch=True, sync_dist=True)
+
+        # Compute and log additional performance metrics.
+        metric_dict = self.performance_metrics(logits=logits)
+        for key in metric_dict:
+            values = metric_dict[key]
+            final_key = 'train_' + key
+            self.log(final_key, metric_dict[key], prog_bar=True if 'f1' in key else False, on_step=True, on_epoch=True, sync_dist=True)
+
+        # Log GPU memory usage at the beginning of the training.
+        if batch_idx == 0:
+            gpu_memory_usage = helper_tools.print_gpu_initialization()
+            self.log(f'gpu_memory_usage', gpu_memory_usage, sync_dist=True)
+        
+        # log CPU memory
+        memory_usage = helper_tools.print_memory_usage()
+        self.log(f'memory_usage', memory_usage, sync_dist=True)
+ 
+        return {'loss': loss}
+
+
+    def validation_step(
+            self,
+            batch: torch.Tensor,
+            batch_idx: any,
+        ) -> dict:
+
+        """
+        `validation_step()`: Validates a single batch of data and computes loss and performance metrics.
+
+        Parameters:
+        - `self`: Reference to the current instance of the model or module.
+        - `batch`: Input data, which might contain text and protein sequences, their corresponding masks, and additional data from both Swiss-Prot and Pfam datasets.
+        - `batch_idx`: Identifier for the current batch.
+
+        Functionality:
+        1. Extracts and processes data from the given batch.
+        2. Computes embeddings for Swiss-Prot and Pfam datasets.
+        3. Concatenates these embeddings to form a unified representation.
+        4. Computes various loss values: inter-modal, intra-modal, and masked language losses for both biomedical texts and protein sequences.
+        5. Logs the computed loss values and other performance metrics, highlighting metrics such as F1-score.
+        6. Collects and appends the joint embeddings of the batch for potential future use.
+
+        Returns:
+        - A dictionary with the total validation loss for the current batch.
+        """
+
+        if isinstance(batch, list):
+            # split the data
+            text_batch, protein_batch, text_mask_batch, protein_mask_batch, \
+            pfam_text_batch, pfam_protein_batch, pfam_text_mask_batch, pfam_protein_mask_batch, \
+            bool_pfam_vector = batch
+
+        
+        # forward pass over the swiss-prot data
+        z_t_swiss, z_p_swiss = self(
+                                  x_t=text_batch,
+                                  x_p=protein_batch,
+                                  compute_masked_logits=False
+        )
+        dist.barrier() # wait till all GPUs catch up...
+     
+        # gather all tensors
+        z_t_swiss_all = self.all_gather(z_t_swiss, sync_grads=True)
+        dist.barrier()    
+        z_p_swiss_all = self.all_gather(z_p_swiss, sync_grads=True)
+
+        # stack the embeddings
+        z_t_swiss_all = z_t_swiss_all.view(-1, z_t_swiss.shape[-1])
+        z_p_swiss_all = z_p_swiss_all.view(-1, z_p_swiss.shape[-1])
+
+        # foward pass over the pfam data
+        z_t_pfam, z_p_pfam = self(
+                                x_t=pfam_text_batch,
+                                x_p=pfam_protein_batch,
+                                compute_masked_logits=False
+        )
+        dist.barrier() # wait till all GPUs catch up...
+        
+        # gather all tensors
+        z_t_pfam_all = self.all_gather(z_t_pfam, sync_grads=True)
+        dist.barrier()
+        z_p_pfam_all = self.all_gather(z_p_pfam, sync_grads=True)
+        
+        # stack the embeddings
+        z_t_pfam_all = z_t_pfam_all.view(-1, z_t_pfam.shape[-1])
+        z_p_pfam_all = z_p_pfam_all.view(-1, z_p_pfam.shape[-1])
+           
+        # concatenate swiss-prot <> pfam embeddings
+        z_t_all = torch.cat((z_t_swiss_all, z_t_pfam_all), dim=0)
+        z_p_all = torch.cat((z_p_swiss_all, z_p_pfam_all), dim=0)
+
+        # compute inter-modal loss values     
+        loss_align, logits = self.model.compute_inter_loss(
+                                            protein_embeddings=z_p_all,
+                                            text_embeddings=z_t_all,
+                                            batch_size=z_p_all.shape[0] // 2
+        )
+
+        # compute intra-modal loss values
+        loss_intra, cosine_similarity = self.model.compute_intra_loss(
+                                            protein_embeddings=z_p_all,
+                                            batch_size=z_p_all.shape[0] // 2
+        )
+
+        # concatenate batch samples
+        all_text_batch = torch.cat((text_batch, pfam_text_batch), dim=0)
+        all_protein_batch = torch.cat((protein_batch, pfam_protein_batch), dim=0)
+        all_text_mask_batch = torch.cat((text_mask_batch, pfam_text_mask_batch), dim=0)
+        all_protein_mask_batch = torch.cat((protein_mask_batch, pfam_protein_mask_batch), dim=0)
+
+        # compute mask language model logits
+        logits_t_mask, logits_s_mask = self(
+                    x_t=all_text_mask_batch,
+                    x_p=all_protein_mask_batch,
+                    compute_masked_logits=True
+        )
+
+        # compute mask language loss for biomedical expert model
+        loss_text_mask = self.model.compute_masked_lang_loss(
+                    logits_masked=logits_t_mask,
+                    targets=all_text_batch,
+                    targets_masked=all_text_mask_batch,
+                    mask_token_id=self.text_tokenizer.mask_token_id
+        )
+
+        # compute mask language loss for protein expert model
+        loss_sequence_mask = self.model.compute_masked_lang_loss(
+                    logits_masked=logits_s_mask,
+                    targets=all_protein_batch,
+                    targets_masked=all_protein_mask_batch,
+                    mask_token_id=self.sequence_tokenizer.mask_idx
+        )
+
+
+        # total loss
+        #loss = loss_align + loss_intra + loss_text_mask + loss_sequence_mask
+        
+        if self.script_args.dataset_type == 'pfam':
+            # Aggregate all computed losses.
+            loss = loss_align + loss_intra + loss_text_mask + loss_sequence_mask
+        
+        elif self.script_args.dataset_type == 'pfam_ablated':
+            # Aggregate all losses besides PFC.
+            loss = loss_align + loss_text_mask + loss_sequence_mask
+        else:
+            # Add an assertion here
+            assert self.script_args.dataset_type in ['pfam', 'pfam_ablated'], "Unexpected dataset_type value"
+            sys.stderr.write("Unexpected dataset_type value\n")
+            sys.exit(1)
+
+     
+        # track loss ...
+        self.log('valid_loss', loss, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('valid_loss_align', loss_align, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('valid_loss_intra', loss_intra, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('valid_loss_text_mask', loss_text_mask, prog_bar=False, on_step=True, on_epoch=True, sync_dist=True)
+        self.log('valid_loss_seq_mask', loss_sequence_mask, prog_bar=False, on_step=True, on_epoch=True, sync_dist=True)
+        # log CPU memory
+        memory_usage = helper_tools.print_memory_usage()
+        self.log(f'memory_usage', memory_usage, sync_dist=True)
+ 
+        # track metrics
+        metric_dict = self.performance_metrics(logits=logits.detach().cpu())
+        for key in metric_dict:
+            values = metric_dict[key]
+            final_key = 'valid_' + key
+            self.log(final_key, metric_dict[key], prog_bar=True if 'f1' in key else False, on_step=True, on_epoch=True, sync_dist=True)
+
+
+        # collect joint embedding
+        #self.val_text_joint_latents.append(z_t_all.detach().cpu())
+        #self.val_seq_joint_latents.append(z_p_all.detach().cpu())
+
+        return {'valid_loss': loss}
+
+
+ #   def on_validation_epoch_end(self):
+ #       print('Enter validation end of epoch analysis...')
+#
+ #       # collect and aggregate outputs from all validation steps
+ #       val_z_t_joint = torch.cat(self.val_text_joint_latents, dim=0)
+ #       val_z_s_joint = torch.cat(self.val_seq_joint_latents, dim=0)
+ #       
+ #       # compute singular values 
+ #       print('Compute singular values...')
+ #       text_log_sigma_k, S_text = self.compute_singular(val_z_t_joint.detach().cpu())
+ #       protein_log_sigma_k, S_protein = self.compute_singular(val_z_s_joint.detach().cpu())
+ #      
+ #       # save image pngs for tracking dimensionality collapse
+ #       self.save_png_to_tensorboard(
+ #               data=text_log_sigma_k.numpy(),
+ #               title='text',
+ #       )
+ #       self.save_png_to_tensorboard(
+ #               data=protein_log_sigma_k.numpy(),
+ #               title='protein'
+ #       )
+ #
+ #       # free memory
+ #       self.val_text_joint_latents.clear()
+ #       self.val_seq_joint_latents.clear()
+ #       
+ #       
+ #       # compute effective rank (RankME):
+ #       print('Compute eranks')
+ #       erank_text = self.compute_effective_rank(sigma_ks=S_text)
+ #       erank_protein = self.compute_effective_rank(sigma_ks=S_protein)
+ #       
+ #       # log erank metrics
+ #       self.log('valid_erank_text', erank_text, sync_dist=True)
+ #       self.log('valid_erank_protein', erank_protein, sync_dist=True)
+
+    def configure_optimizers(self,):
+
+        params = [
+                {"params": self.model.protein_encoder.parameters(), "lr": self.script_args.protein_encoder_lr},
+                {"params": self.model.text_encoder.parameters(), "lr": self.script_args.text_encoder_lr},
+                {"params": itertools.chain(
+                    self.model.protein_projection.parameters(),
+                    self.model.text_projection.parameters()
+                    ),
+                "lr": self.script_args.head_lr,
+                "weight_decay": self.script_args.weight_decay}
+        ]
+
+        optimizer = torch.optim.AdamW(params, weight_decay=self.script_args.weight_decay)
+
+        return {
+                "optimizer": optimizer,
+        }
+    
+    @torch.no_grad()
+    def compute_class_metrics(
+            self,
+            outputs: torch.Tensor,
+            targets: torch.Tensor,
+            source: str
+        ) -> dict:
+
+        # convert torch tensors to numpy array
+        outputs_np = outputs.numpy()
+        targets_np = targets.numpy()
+
+        # compute the metrics
+        accuracy = accuracy_score(targets_np, outputs_np.round())
+        precision = precision_score(targets_np, outputs_np.round(), average='micro')
+        recall = recall_score(targets_np, outputs_np.round(), average='micro')
+        f1 = f1_score(targets_np, outputs_np.round(), average='micro')
+
+        return {
+                f'{source}_accuracy': accuracy,
+                f'{source}_precision': precision,
+                f'{source}_recall': recall,
+                f'{source}_f1': f1
+        }
+
+    @torch.no_grad()
+    def performance_metrics(self, logits: torch.Tensor) -> tuple:
+
+        logits = logits.cpu().float()
+
+        # get probs
+        p_text = F.softmax(logits, dim=-1) # prob of a given text captions aligning well with seq. pairs
+        p_seq = F.softmax(logits.T, dim=-1) # prob of a given seq aligning well with text pairs
+        p_tot = (p_seq + p_text) / 2 # total prob
+
+        # get class labels
+        y_pred_text = torch.argmax(p_text, dim=-1)
+        y_pred_seq = torch.argmax(p_seq, dim=-1)
+        y_pred = torch.argmax(p_tot, dim=-1)
+        y_true = torch.arange(y_pred_text.shape[0])
+        
+        # compute class metrics
+        text_metrics = self.compute_class_metrics(
+                                    outputs=y_pred_text,
+                                    targets=y_true,
+                                    source='text'
+        )
+        seq_metrics = self.compute_class_metrics(
+                                    outputs=y_pred_seq,
+                                    targets=y_true,
+                                    source='seq'
+        )
+        total_metrics = self.compute_class_metrics(
+                                    outputs=y_pred,
+                                    targets=y_true,
+                                    source='total'
+        )
+
+        # combine dicts into one
+        combined_dict = {}
+        combined_dict.update(text_metrics)
+        combined_dict.update(seq_metrics)
+        combined_dict.update(total_metrics)
+
+        return combined_dict
+    
+    @torch.no_grad()
+    def compute_singular(self, inputs: torch.Tensor) -> (
+            torch.Tensor,
+            torch.Tensor
+        ):
+
+        # goal of this function: track for dimensionality collapse
+        # inputs dim: (batch_size, emb_dim)
+
+        mean_inputs = torch.mean(inputs, dim=0) # average over batch dimension
+        norm_inputs = inputs - mean_inputs # normalize vectors
+       
+       # compute correlation matrix  #TODO: double check work...
+        C = torch.zeros((norm_inputs.shape[-1], norm_inputs.shape[-1]))
+        for sample_idx in range(norm_inputs.shape[0]):
+            norm_vector = norm_inputs[sample_idx, :].unsqueeze(0)
+            C += norm_vector.T @ norm_vector
+        C *= 1/norm_vector.shape[0]
+
+        _, S, _ = torch.linalg.svd(C, full_matrices=False)
+        
+        # return singular value indexes 
+        log_sigma_k, _ = torch.sort(torch.log(S), descending=True)
+        return (
+                log_sigma_k,
+                S
+        )
+        
+    def compute_effective_rank(self, sigma_ks: torch.Tensor) -> torch.Tensor:
+        """
+        references:
+            - Roy et al. The effective rank: a measure of effective dimensionality
+            - Garrido et al. RankMe: Assessing the Downstream Performnace of Pretrained SS Reps by their Rank.
+        """
+        # sort the singular values
+        sigma_ks, _ = torch.sort(sigma_ks, descending=True)
+        
+        # copute L1 norm for sing values.
+        l1_norm_sigma = torch.norm(sigma_ks, p=1)
+        
+        # compute singular value distribution
+        p_k = sigma_ks / l1_norm_sigma  + torch.finfo(torch.float).eps
+        
+        # compute Shannon entropy
+        entropy = - torch.sum(p_k * torch.log(p_k))
+    
+        # get effective rank (RankME):
+        erank = torch.exp(entropy)
+
+        return erank
+
+    def save_png_to_tensorboard(
+            self,
+            data: np.single,
+            title: str,
+            x_axis_label: str='Singular Value Rank Index',
+            y_axis_label: str='Log of singular values',
+            ):
+    
+        current_epoch = self.trainer.current_epoch
+        
+        # Plot the line
+        fig, ax = plt.subplots(dpi=300)
+        ax.plot(data)
+        ax.set_xlabel(x_axis_label)
+        ax.set_ylabel(y_axis_label)
+        ax.set_title(title)
+        ax.set_ylim([-25,3])
+
+        # Log the plot in TensorBoard 
+        self.logger.experiment.add_figure(f'{title}_SingularValues_{current_epoch}', fig, current_epoch)
+        
+        # Close the figure to free up memory
+        plt.close(fig)
+
+    def predict_step(
+            self,
+            batch: torch.Tensor,
+            batch_idx: torch.Tensor,
+            dataloder_idx: bool=False
+        ) -> (  
+                torch.Tensor,
+                torch.Tensor
+        ):
+
+
+        if isinstance(batch, list):
+            # mean loss
+            text_batch, protein_batch = batch
+            outputs = self(
+                    x_t=text_batch,
+                    x_p=protein_batch,
+                    compute_masked_logits=False
+            )
+        
+        z_t_joint, z_p_joint = outputs
+
+        self.predict_text_joint_latents.append(z_t_joint.detach().cpu())
+        self.predict_seq_joint_latents.append(z_p_joint.detach().cpu())
+
+        return outputs
+
+    def on_predict_epoch_end(self, outputs=None):
+
+        self.predict_text_joint_latents = torch.cat(self.predict_text_joint_latents).cpu()
+        self.predict_seq_joint_latents = torch.cat(self.predict_seq_joint_latents).cpu()
+
+
+##########################
+# Facilitator PL wrapper #
+##########################
+
+class PL_Facilitator(pl.LightningModule):
+
+    def __init__(
+            self, 
+            args: any
+    ):
+
+        super().__init__()
+
+        # arguments
+        self.args = args
+
+        # model 
+        self.model = mod.Facilitator(
+                    in_dim=self.args.emb_dim,
+                    hid_dim=self.args.hid_dim,
+                    out_dim=self.args.emb_dim,
+                    dropout=self.args.dropout
+        )
+        
+        self.text_to_protein_joint_embeddings = []
+
+    def forward(
+            self,
+            z_t: torch.Tensor,
+    ) -> torch.Tensor:
+
+        # reconfigure z_t to z_p (additional alignment)
+        z_t_to_p = self.model(z_t)
+
+        return z_t_to_p
+
+    
+
+    def training_step(self, batch: torch.Tensor, batch_id: any) -> dict:
+
+        # check if the batch is a list and split data if so 
+        if isinstance(batch, list):
+            text_embeddings, protein_embeddings = batch
+
+        # forward pass with the model
+        z_t_to_p = self(z_t=text_embeddings)
+
+        # compute loss
+        loss = self.model.compute_loss(
+                output=z_t_to_p,
+                target=protein_embeddings,
+                loss_option=self.args.loss_type
+        )
+    
+        # log the total loss
+        self.log('train_loss', loss, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        return {'loss': loss}
+
+ 
+    def validation_step(self, batch: torch.Tensor, batch_id: any) -> dict:
+
+        # check if the batch is a list and split data if so 
+        if isinstance(batch, list):
+            text_embeddings, protein_embeddings = batch
+
+        # forward pass with the model
+        z_t_to_p = self(z_t=text_embeddings)
+    
+        # compute loss
+        loss = self.model.compute_loss(
+                output=z_t_to_p,
+                target=protein_embeddings,
+                loss_option=self.args.loss_type
+        )
+    
+        # log the total loss
+        self.log('valid_loss', loss, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        return {'loss': loss}
+
+    
+    def configure_optimizers(self,):
+
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.args.lr, weight_decay=self.args.weight_decay)
+    
+        return {
+                "optimizer": optimizer
+        }
+   
+
+    def predict_step(self, batch: torch.Tensor, batch_idx: int, dataloader_idx: int = None) -> torch.Tensor:
+        """
+        Defines a single prediction (inference) step.
+        """
+
+        # Unpack the batch if it comes in a list format.
+        # Here, we only take text embeddings for prediction as an example.
+        if isinstance(batch, list):
+            text_embeddings, _ = batch  # We ignore the second element (protein_embeddings)
+        else:
+            text_embeddings = batch
+
+        # Perform forward pass to get transformed text embeddings (z_t_to_p)
+        z_t_to_p = self(z_t=text_embeddings)
+        self.text_to_protein_joint_embeddings.append(z_t_to_p.detach().cpu())
+
+        return z_t_to_p
+
+    def on_predict_epoch_end(self, outputs=None):
+        
+        self.text_to_protein_joint_embeddings = torch.cat(self.text_to_protein_joint_embeddings).cpu()

Stage1_source/helper_funcs.py

@@ -0,0 +1,37 @@
+from pynvml import *
+import psutil
+
+"""
+To track memory allocation, let's take advantage of the nvidia-ml-py3 package and GPU memory allocation from python.
+
+ref: https://huggingface.co/docs/transformers/v4.20.1/en/perf_train_gpu_one
+"""
+
+
+def print_gpu_initialization():
+    nvmlInit()
+    handle = nvmlDeviceGetHandleByIndex(0)
+    info = nvmlDeviceGetMemoryInfo(handle)
+    print(f"GPU memory occupied: {info.used//1024**2} MB.")
+    return info.used // 1024**2
+
+
+def print_summary(result):
+    print(f"Time: {result.metrics['train_runtime']:.2f}")
+    print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
+    print_gpu_utilization()
+
+
+def print_memory_usage():
+    process = psutil.Process(os.getpid())
+    memory_in_bytes = process.memory_info().rss
+    memory_in_megabytes = memory_in_bytes / (1024 ** 2)
+    #print(f"Memory used by this script: {memory_in_megabytes:.2f} MB")
+
+    return memory_in_megabytes
+
+
+
+
+
+

Stage1_source/model.py

@@ -0,0 +1,556 @@
+import os
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModel, BertTokenizer, BertForMaskedLM
+import torch.distributed as dist
+import esm
+from torch.nn.utils.weight_norm import weight_norm
+
+
+"""
+functions and classes adapted from the following:
+    1. https://keras.io/examples/vision/nl_image_search/
+    2. https://colab.research.google.com/drive/1hYHb0FTdKQCXZs3qCwVZnSuVGrZU2Z1w?usp=sharing
+"""
+
+class ProteinEncoder(nn.Module):
+    """
+    Encoder for protein sequence to a fixed size vector --> z_s
+    """
+    
+    def __init__(self, args: any):
+        super().__init__()
+
+        #self.script_args = args
+        self.seq_model_path = args.seq_model_path
+        self.pretrained = args.pretrained_seq
+        self.trainable = args.trainable_seq
+        self.n_layers_to_finetune = args.pLM_n_layers_to_finetune
+        self.rep_layer = args.rep_layer
+        self.model, self.alphabet = self.get_ESM_model() # get model and alphabet (ESM)
+        
+        for p in self.model.parameters():
+            if self.trainable and self.n_layers_to_finetune == 0:
+                p.required_grad = True
+            else:
+                p.requires_grad = False
+        
+        # Make the last n_layers_to_finetune layers trainable
+        if self.trainable and self.n_layers_to_finetune != 0:
+            for layer in self.model.layers[-self.n_layers_to_finetune:]:
+                for p in layer.parameters():
+                    p.requires_grad = True
+
+        # Use the [CLS] token hidden representation as the sentence's embedding
+        # for the downstream latent alignment.
+        self.target_token_idx = 0
+
+    def get_ESM_model(self):
+
+        return esm.pretrained.load_model_and_alphabet(
+                os.path.expanduser(
+                    self.seq_model_path
+                )
+        )
+    
+    def forward(self, x_s: torch.Tensor, compute_logits: bool=False):
+        # drop channel depth
+        x_s = x_s.squeeze(1)
+
+        outputs = self.model(
+                x_s,
+                repr_layers=[self.rep_layer],
+                return_contacts=False
+        )
+        
+        # mask langauge model objective 
+        if compute_logits:
+            logits = outputs['logits']
+            return logits
+       
+        # fine-tuning cls token for protein sequence alignment with biomedical text 
+        cls_hidden = outputs['representations'][self.rep_layer][:,self.target_token_idx,:]
+        return cls_hidden
+    
+class TextEncoder(nn.Module):
+
+    """
+    Encoder for protein's natural text to a fixed size vector --> z_t
+    """
+
+    def __init__(self, args: any):
+        super().__init__()
+ 
+        self.model_name = args.text_model_path
+        self.pretrained = args.pretrained_text
+        self.trainable = args.trainable_text
+        self.n_layers_to_finetune = args.bLM_n_layers_to_finetune 
+        self.tokenizer = AutoTokenizer.from_pretrained(args.text_model_path)
+
+        if self.pretrained:
+            #self.model = AutoModel.from_pretrained(self.model_name)
+            self.model = BertForMaskedLM.from_pretrained(self.model_name)
+
+        else:
+            #self.model = AutoModel.from_config(self.model_name)
+            self.model = BertForMaskedLM.from_config(self.model_name)
+
+        for p in self.model.parameters():
+            if self.trainable and self.n_layers_to_finetune == 0:
+                p.required_grad = True
+            else:
+                p.requires_grad = False
+       
+        # Make the last n_layers_to_finetune layers trainable
+        if self.trainable and self.n_layers_to_finetune != 0:
+            for layer in self.model.bert.encoder.layer[-self.n_layers_to_finetune:]:
+                for p in layer.parameters():
+                    p.requires_grad = True
+        
+        # Use the [CLS] token hidden representation as the sentence's embedding
+        # for the downstream latent alignment.
+        self.target_token_idx = 0
+
+    def forward(self, inputs: torch.Tensor, compute_logits: bool=False) -> torch.Tensor:
+        # drop channel depth
+        inputs = inputs.squeeze(1)
+        
+        if compute_logits:
+            # compute the masked language model logits
+            #sequence_output = outputs.last_hidden_state
+            outputs = self.model(inputs)
+            logits = outputs.logits
+            return logits
+        
+        else:
+            outputs = self.model(inputs, output_hidden_states=True)
+            # use the token representations...
+            last_hidden_state = outputs.hidden_states[-1]
+            return last_hidden_state[:, self.target_token_idx, :] # return [cls] token 
+
+
+
+class ProjectionHead(nn.Module):
+    """
+    g(.) which maps z_t --> h_t or z_s --> h_s
+    
+    Note: h is the joint embedding representation, h_t
+    is the joint embedding for the text caption, and
+    h_s is the joint embedding for the protein sequence.
+    """
+
+    def __init__(self, embedding_dim: int, args: any):
+
+        super().__init__()
+        self.projection_dim = args.proj_embedding_dim
+        self.dropout = args.dropout
+        self.embedding_dim = embedding_dim
+
+        # model graph
+        self.projection = nn.Linear(self.embedding_dim, self.projection_dim)
+        self.gelu = nn.GELU()
+        self.fc = nn.Linear(self.projection_dim, self.projection_dim)
+        self.dropout = nn.Dropout(self.dropout)
+        self.layer_norm = nn.LayerNorm(self.projection_dim)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+
+        projection = self.projection(z)
+        h = self.gelu(projection)
+        h = self.fc(h)
+        h = self.dropout(h)
+        h = h + projection
+        h = self.layer_norm(h)
+        return h
+
+
+
+
+#####################
+# Pfam architecture #
+#####################
+
+
+
+class pfam_PEN_CL(nn.Module):
+
+    """
+    Protein Embeddings with Natural lanauge using Constrastive Learing (PEN-CL) while including pfam constrastive learning.
+    """
+
+    def __init__(self, args: any):
+
+        super().__init__()
+
+        self.protein_embedding = args.protein_encoder_embedding
+        self.text_embedding = args.text_encoder_embedding
+        self.temperature = args.temperature
+
+        # protein sequence expert
+        self.protein_encoder = ProteinEncoder(args=args)
+        # natural text expert
+        self.text_encoder = TextEncoder(args=args)
+        
+        # projection heads g_seq( . ) --> joint embedding space
+        self.protein_projection = ProjectionHead(
+                embedding_dim=self.protein_embedding,
+                args=args
+        )
+
+        # projection heads g_text( . ) --> joint embedding space
+        self.text_projection = ProjectionHead(
+                embedding_dim=self.text_embedding,
+                args=args
+        )
+
+    def forward(
+            self,
+            x_t: torch.Tensor,
+            x_s: torch.Tensor,
+            compute_masked_logits: bool=False
+        ) -> dict:
+
+        if compute_masked_logits:
+            # forward pass for computing logits for masked langauge objective
+            protein_logits = self.protein_encoder(x_s, compute_logits=True)
+            text_logits = self.text_encoder(x_t, compute_logits=True)
+
+            return {
+                    'text_masked_logits': text_logits,
+                    'protein_masked_logits': protein_logits
+            }
+
+        else:
+            # split the tuple into 2 dicts... 
+            # getting protein sequence and text inputs ...
+            z_t = self.text_encoder(x_t, compute_logits=False)
+            z_s = self.protein_encoder(x_s, compute_logits=False)
+
+            # "joint" sequence and text embedding (with same dimension)
+            z_t_joint = self.text_projection(z_t)
+            z_s_joint = self.protein_projection(z_s)
+
+            return {
+                    'text_joint_latent': z_t_joint,
+                    'seq_joint_latent': z_s_joint,
+            }
+
+    def compute_inter_loss(
+        self,
+        protein_embeddings: torch.Tensor,
+        text_embeddings: torch.Tensor,
+        batch_size: int
+    ) -> (
+            torch.Tensor,
+            torch.Tensor
+            ):
+            
+            """
+            Compute the inter-modal contrastive InfoNCE loss between protein and text embeddings.
+
+            Parameters:
+            - protein_embeddings: A tensor representing the embeddings of the protein sequences.
+            - text_embeddings: A tensor representing the embeddings of the text descriptions.
+            - batch_size: The number of samples in the batch.
+
+            Steps:
+            1. Generate a masking matrix to identify off-diagonal elements.
+            2. Compute cosine similarities (i.e., logits) between text and protein embeddings.
+            3. Compute self-similarities for both protein and text embeddings.
+            4. Mask off-diagonal elements between swiss-prot and pfam in the similarity matrices.
+            5. Define ground truth by averaging the masked protein and text similarity matrices.
+            6. Compute the contrastive loss for the protein and text embeddings using the ground truth.
+
+            Returns:
+            - Mean contrastive loss for the given batch of protein and text embeddings.
+            - The logits (cosine similarity matrix between text and protein embeddings).
+
+            Note: This function assumes a specific structure in the input batches, where corresponding positive samples 
+            in the protein and text embeddings are arranged in a particular way, allowing for masking and contrastive loss calculation.
+            """
+            
+            # get off-diagonal masking matrix
+            mask = torch.zeros((2*batch_size, 2*batch_size))
+            # mask the bottom left quadrant diagonal
+            mask[batch_size:, :batch_size] = torch.eye(batch_size)
+            # mask the top right quadrant
+            mask[:batch_size, batch_size:] = torch.eye(batch_size)
+            # convert to correct device and convert to boolean
+            mask = mask.to(protein_embeddings.device).bool()
+
+            # matrix multiplication between model embeddings
+            logits = (text_embeddings @ protein_embeddings.T) / self.temperature
+            protein_similarity = protein_embeddings @ protein_embeddings.T
+            text_similarity = text_embeddings @ text_embeddings.T
+
+            # mask the off-diagonal between swiss-prot and pfam
+            mask_protein_similarity = self.set_inf(protein_similarity, mask)
+            mask_text_similarity = self.set_inf(text_similarity, mask)
+            mask_logits = self.set_inf(logits, mask)
+
+            # ground truth
+            targets = F.softmax(
+                (mask_protein_similarity + mask_text_similarity) / (2 * self.temperature), dim=-1
+            )
+
+            # compute loss
+            text_loss = self.cross_entropy(mask_logits, targets, reduction='none')
+            protein_loss = self.cross_entropy(mask_logits.T, targets.T, reduction='none')
+            loss = (protein_loss + text_loss) / 2.0
+
+            return (
+                loss.mean(),
+                mask_logits.detach().cpu()
+            )
+
+
+    def compute_intra_loss(  
+            self,
+            protein_embeddings,
+            batch_size
+        ) -> (
+                torch.Tensor,
+                torch.Tensor,
+        ):
+        """
+        Compute the intra-modal contrastive InfoNCE loss for protein embeddings.
+
+        Parameters:
+        - protein_embeddings: A tensor representing the embeddings of the protein sequences.
+        - batch_size: Batch size used for training.
+
+        Steps:
+        1. Normalize the protein embeddings using L2 normalization.
+        2. Compute the cosine similarity between the normalized embeddings.
+        3. Mask the diagonal of the cosine similarity matrix to avoid using a protein's similarity with itself.
+        4. Define positive examples by rolling the mask. The positive example for a given protein embedding is determined by an embedding half the batch size away.
+        5. Compute the InfoNCE loss using the masked cosine similarity matrix.
+
+        Returns:
+        - Mean InfoNCE loss for the given batch of protein embeddings.
+        - The cosine similarity matrix.
+
+        Note: The underlying assumption is that in each batch, corresponding positive samples for a given protein embedding 
+        lie half the batch size away. The function computes the negative log likelihood loss between these positive samples 
+        and the entire batch.
+        """
+        
+        # l2 normalization
+        #norm_protein_embeddings = F.normalize(protein_embeddings, p=2, dim=1)
+        norm_protein_embeddings = protein_embeddings
+
+        # cosine similarity
+        cosine_similarity = (norm_protein_embeddings @ norm_protein_embeddings.T) / self.temperature
+        
+        # mask cosine similarity matrix
+        sample_size = protein_embeddings.shape[0]
+        mask = torch.eye(sample_size, device=cosine_similarity.device, dtype=torch.bool)
+        #cosine_similarity.masked_fill_(mask, float(-9e15))
+        cosine_similarity = self.set_inf(cosine_similarity, mask)
+        
+        # Find positive example -> batch_size //2 away from the original example (swiss-prot<>pfam)
+        pos_mask = mask.roll(shifts=mask.shape[0]//2, dims=0)
+
+        # InfoNCE loss
+        nll = -cosine_similarity[pos_mask] + torch.logsumexp(cosine_similarity, dim=-1)
+        
+        return (
+            nll.mean(),
+            cosine_similarity.cpu(),
+        )
+
+    def set_inf(
+            self,
+            tensor: torch.Tensor,
+            mask: torch.Tensor
+        ) -> torch.Tensor:
+        # Determine replacement value based on tensor dtype
+        if tensor.dtype == torch.float32:
+            replace_value = -9e15
+        elif tensor.dtype == torch.float16:
+            replace_value = -1e4
+        else:
+            raise ValueError("Unsupported tensor dtype for this operation.")
+
+        # Use masked_fill_ to replace positions in tensor where mask is True with the specified value
+        tensor.masked_fill_(mask, replace_value)
+
+        return tensor
+
+    def cross_entropy(
+            self,
+            preds: torch.Tensor,
+            targets: torch.Tensor,
+            reduction: str='none'
+        ) -> torch.Tensor:
+
+        # compute categorical cross entropy
+        log_softmax = nn.LogSoftmax(dim=-1)
+        loss = (-targets * log_softmax(preds)).sum(1)
+
+        if reduction == 'none':
+            return loss
+        elif reduction == 'mean':
+            return loss.mean()
+        else:
+            assert False, print('Choose either "none" or "mean" for reduction argument')
+
+    def compute_masked_lang_loss(
+            self,
+            logits_masked: torch.Tensor,
+            targets: torch.Tensor,
+            targets_masked: torch.Tensor,
+            mask_token_id: torch.Tensor
+        ) -> torch.Tensor:
+        
+        """
+        Compute the masked language model loss for BERT-like architectures.
+
+        Given a batch of logits predicted for masked positions and their corresponding target tokens, this function 
+        computes the cross-entropy loss between the predicted logits and the true labels, but only for positions 
+        that have been masked in the input.
+
+        Parameters:
+        - logits_masked: Predicted token logits for masked positions from the model.
+                         Shape: (batch_size, seq_len, vocab_size).
+        - targets: True token IDs for each position in the input sequence.
+                   Shape: (batch_size, seq_len).
+        - targets_masked: Token IDs for the input sequence, including masked positions.
+                          Shape: (batch_size, seq_len).
+        - mask_token_id: The ID corresponding to the [MASK] token in the vocabulary.
+
+        Steps:
+        1. Compute the cross-entropy loss between predicted logits and true labels across all positions.
+        2. For each sample in the batch, locate the positions that were masked.
+        3. Extract the loss values corresponding to these masked positions.
+        4. Compute and return the mean of these extracted loss values across the batch.
+
+        Returns:
+        - Mean cross-entropy loss for masked positions across the batch.
+
+        Note: This function focuses exclusively on masked positions in the input, as is typical for the MLM objective 
+        in BERT-like models. It disregards unmasked positions.
+        """
+
+        # compute the masked langauge objective loss for masked logits
+        loss_func = nn.CrossEntropyLoss(reduction='none')
+        loss_mask = loss_func(
+                            logits_masked.permute(0, 2, 1), # (batch_size, vocab_size, seq_len)
+                            targets.squeeze(1) # (batch_size, seq_len)
+        )
+
+        # list to append loss values
+        batch_loss = []
+
+        for ii, target_mask_sample in enumerate(targets_masked):
+            
+            # locate mask positions 
+            masked_positions = (target_mask_sample == mask_token_id).tolist()
+            # extract the loss values at those masked positions
+            loss_mask_sample = loss_mask[ii][masked_positions]
+            
+            # append mean loss value for a given batch sample
+            if loss_mask_sample.numel() > 0:
+                batch_loss.append(torch.mean(loss_mask_sample).unsqueeze(0))
+        
+        if len(loss_mask_sample) > 0:
+            loss_mask_mean = torch.mean(torch.cat(batch_loss))
+        else:
+            # handle the case where there are no masked positions in any sample 
+            loss_mask_mean = torch.tensor(0.0, device=logits_masked.device)
+
+        return loss_mask_mean
+
+
+###############
+# Facilitator #
+###############
+
+
+class Facilitator(nn.Module):
+
+    def __init__(self,
+                 in_dim: int,  # Input dimension
+                 hid_dim: int,  # Hidden layer dimension
+                 out_dim: int,  # Output dimension
+                 dropout: float = 0.  # Dropout rate
+                 ):
+        super().__init__()
+
+        # Main neural network structure
+        self.main = nn.Sequential(
+            weight_norm(nn.Linear(in_dim, hid_dim), dim=None),  # Weight-normalized linear layer
+            nn.GELU(),  # GELU activation function
+            nn.Dropout(dropout, inplace=True),  # Dropout layer
+            weight_norm(nn.Linear(hid_dim, out_dim), dim=None)  # Weight-normalized output layer
+        )
+
+    def forward(self, x):
+        # Forward pass through the network
+        return self.main(x)
+
+    def compute_loss(self, output: torch.Tensor, target: torch.Tensor, loss_option='MSE') -> torch.Tensor:
+        # Compute loss based on the chosen loss_option ('MSE' or 'MMD')
+        if loss_option == 'MSE':
+            return Facilitator.compute_MSE(output, target)
+        elif loss_option == 'MMD':
+            return Facilitator.compute_mmd(output, target)
+        else:
+            return ValueError("Invalid loss option")
+    
+    @staticmethod
+    def compute_MSE(output, target):
+        # Compute Mean Squared Error between output and target
+        mse_loss = nn.MSELoss()
+        loss = mse_loss(output, target)
+        return loss
+
+    @staticmethod
+    def compute_kernel(
+            x: torch.FloatTensor,
+            y: torch.FloatTensor
+        ) -> torch.FloatTensor:
+        """
+        Compute the Gaussian RBF kernel between tensors x and y
+        """
+
+        # Get the sizes of each mini-batch
+        x_size, y_size = x.shape[0], y.shape[0]
+
+        # Dimension based on z size
+        dim = x.shape[1]
+
+        x = x.view(x_size, 1, dim)
+        y = y.view(1, y_size, dim)
+
+        x_core = x.expand(x_size, y_size, dim)
+        y_core = y.expand(x_size, y_size, dim)
+
+        # Gaussian RBF kernel computation
+        return torch.exp(-(x_core - y_core).pow(2).mean(2) / dim)
+
+    @staticmethod
+    def compute_mmd(
+            x: torch.FloatTensor,
+            y: torch.FloatTensor
+        ) -> torch.FloatTensor:
+        """
+        Compute the Maximum Mean Discrepancy (MMD) between two distributions.
+        Args:
+            x: Samples from first distribution (z_t_to_p ~ q(z_p))
+            y: Samples from second distribution (z_p ~ p(z_p))
+        Returns:
+            MMD_loss: The MMD loss between the sampled distributions
+        """
+
+        x_kernel = Facilitator.compute_kernel(x, x)
+        y_kernel = Facilitator.compute_kernel(y, y)
+        xy_kernel = Facilitator.compute_kernel(x, y)
+
+        # Calculate MMD loss
+        return x_kernel.mean() + y_kernel.mean() - 2 * xy_kernel.mean()
+
+

Stage1_source/preprocess.py

@@ -0,0 +1,410 @@
+import torch
+from torch.utils.data import random_split, Dataset, DataLoader, Subset, ConcatDataset
+import pandas as pd
+import random
+import ast
+import dask.dataframe as dd
+import os
+from sklearn.model_selection import train_test_split
+from pytorch_lightning import LightningDataModule
+from tqdm import tqdm
+import gc
+import psutil
+import time
+import copy
+
+import esm
+from esm import pretrained
+from transformers import AutoTokenizer, AutoModel
+
+
+########################################
+# Dataset iterator with masking tokens #
+########################################
+
+class TextSeqPairing_Dataset(Dataset):
+
+    def __init__(self, args: any, df: pd.Series):
+
+        # dataframe
+        self.df = df
+        self.length = self.df.shape[0]
+        self.df_column_names = self.df.columns.tolist()
+        self.protein_sequence_list = self.df[args.sequence_keyword].tolist()
+        self.text_captions_list = self.df['[final]text_caption'].tolist()
+        self.accession_id_list = self.df[args.id_keyword].tolist()
+
+        # parameters
+        self.text_max_length = args.text_max_length # max BERT sequence tokenization length
+        self.seq_max_length = 1024 # max ESM model
+
+        # tokenizers
+        self.text_tokenizer = AutoTokenizer.from_pretrained(args.text_model_path) # for text encoder
+        _, self.sequence_tokenizer = pretrained.load_model_and_alphabet(args.seq_model_path) # for protein encoder
+
+    def caption_tokenizer(self, batch_captions: list) -> dict:
+        
+        # transform input text tokens
+        text_inputs = self.text_tokenizer.batch_encode_plus(
+                            batch_captions,
+                            truncation=True,
+                            max_length=self.text_max_length,
+                            padding='max_length',
+                            return_tensors='pt',
+                            return_attention_mask=True,
+                            return_token_type_ids=False
+        )
+        
+        # track the original natural language captions
+        text_inputs['orig_captions'] = batch_captions
+
+        return text_inputs
+    
+    def protein_tokenizer(self, batch_sequences: list) -> dict:
+        
+        # perpare data for ESM
+        batch_converter = self.sequence_tokenizer.get_batch_converter()
+        batch_labels, batch_str, batch_tokens = batch_converter(batch_sequences)
+        
+        # pad sequences
+        batch_tokens = torch.cat((
+            batch_tokens,
+            torch.ones((1,1024-batch_tokens.shape[1])),
+            ), dim=-1
+        )
+
+        sequence_inputs = {
+            'protein_sequence_labels': batch_labels, # UniProtKB id
+            'protein_sequence_str': batch_str, # original protein sequence (in amino acids)
+            'protein_sequence_tokens': batch_tokens.long() # training data
+        }
+
+        return sequence_inputs
+    
+    
+    def __getitem__(self, idx: torch.Tensor) -> (
+            dict,
+            dict
+        ):
+        
+        protein_sequence = self.protein_sequence_list[idx]
+        text_captions = self.text_captions_list[idx]
+        accession_id = self.accession_id_list[idx]
+
+        # prepare protein sequence in ESM format (e.g. tuple: (header, sequence)):
+        batch_sequences = [
+            (accession_id, protein_sequence)
+        ]
+        
+        text_data = self.caption_tokenizer(batch_captions=[text_captions])
+        protein_data = self.protein_tokenizer(batch_sequences=batch_sequences)
+ 
+        return (
+                text_data['input_ids'],
+                protein_data['protein_sequence_tokens']
+        )
+
+    def __len__(self):
+        return self.length
+
+
+######################
+# Default DataModule #
+######################
+
+
+class Default_DataModule(LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+
+        # construct dataset iterator
+        dataset_options = {
+                'default': TextSeqPairing_Dataset,
+                'masked': MaskTextSeqPairing_Dataset,
+                'pfam': Pfam_TextSeqPairing_Dataset,
+                'pfam_ablated': Pfam_TextSeqPairing_Dataset
+        }
+
+        self.dataset_class = dataset_options.get(args.dataset_type, TextSeqPairing_Dataset)
+        
+    def prepare_data(self):
+        pass
+
+    def setup(self, stage=None):
+        
+        if self.trainer is not None:
+            print(f"Number of GPUs: {self.trainer.world_size}")
+            print(f"Current GPU index: {self.trainer.local_rank}")
+
+        # Load Swiss-Prot data
+        df = self.load_swiss_prot()
+        
+        # Split the dataframe into train and valid sets
+        train_df, valid_df = train_test_split(
+            df,
+            test_size=self.args.valid_size,
+            random_state=self.args.seed
+        )
+ 
+        print(f"Available memory after pfam_df: {check_available_memory()} GB")
+
+        # Define datasets and dataloaders
+        self.train_dataset = self.dataset_class(args=self.args, df=train_df)
+        self.valid_dataset = self.dataset_class(args=self.args, df=valid_df)
+
+    def load_swiss_prot(self) -> pd.Series:
+        # Load and preprocess data (called on each GPU/TPU in DDP)
+        print(f'Load Swiss-Prot data...')
+
+        # Load Swiss-Prot data
+        df = pd.read_csv(os.path.expanduser(self.args.data_path))
+        df = df[df['protein_sequence'].apply(lambda seq: len(seq) <= 1022)]
+
+        return df
+
+    def train_dataloader(self):
+        return DataLoader(
+                self.train_dataset,
+                batch_size=self.args.batch_size,
+                num_workers=self.args.num_workers,
+                shuffle=True,
+                pin_memory=True
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+                self.valid_dataset,
+                batch_size=self.args.batch_size,
+                num_workers=self.args.num_workers,
+                pin_memory=True
+        )
+
+    def test_dataloader(self):
+        # Define test dataloader if needed
+        pass
+
+
+
+################################
+# Facilitator Dataset Iterator #
+################################
+
+
+class Facilitator_Dataset(Dataset):
+
+    def __init__(self, args: any, dataset: dict):
+
+        # Determine the device based on the number of GPUs
+        device = 'cuda' if args.num_gpus >= 1 else 'cpu'
+
+        # Check if text_embeddings is a list and convert to a tensor
+        if isinstance(dataset['text_embedding'], list):
+            # Convert list elements to tensors if they are not already
+            text_emb_tensors = [torch.tensor(emb).to(device) if not isinstance(emb, torch.Tensor) else emb.to(device) for emb in dataset['text_embedding']]
+            # Stack the list of tensors
+            self.text_embeddings = torch.stack(text_emb_tensors)
+        else:
+            self.text_embeddings = dataset['text_embedding'].to(device)
+
+        # Check if protein_embeddings is a list and convert to a tensor
+        if isinstance(dataset['protein_embedding'], list):
+            # Convert list elements to tensors if they are not already
+            protein_emb_tensors = [torch.tensor(emb).to(device) if not isinstance(emb, torch.Tensor) else emb.to(device) for emb in dataset['protein_embedding']]
+            # Stack the list of tensors
+            self.protein_embeddings = torch.stack(protein_emb_tensors)
+        else:
+            self.protein_embeddings = dataset['protein_embedding'].to(device)
+
+
+    def __getitem__(self, idx: torch.Tensor) -> (
+            torch.Tensor,
+            torch.Tensor
+        ):
+
+
+        z_t = self.text_embeddings[idx]
+        z_p = self.protein_embeddings[idx] 
+
+        return (
+                z_t,
+                z_p
+        )
+
+
+    def __len__(self):
+        return len(self.text_embeddings)
+
+###########################
+# Facilitator Data Module #
+###########################
+
+
+
+class Facilitator_DataModule(LightningDataModule):
+    def __init__(self, args):
+        super().__init__()
+        
+        self.args = args
+       
+        self.OOD_pfam_labels = [
+                'PF18369', # Polyketide synthase dimerisation element domain
+                'PF04680', # Opioid growth factor receptor repeat
+                'PF17988', # VEGFR-2 Transmembrane domain
+                'PF12325', # TATA element modulatory factor 1 TATA binding
+                'PF03272', # Putative mucin or carbohydrate-binding module
+                'PF03938', # Outer membrane protein (OmpH-like)
+                'PF17724', # Family of unknown function (DUF5568)
+                'PF10696', # Protein of unknown function
+                'PF11968', # 25S rRNA (adenine(2142)-N(1))-methyltransferase, Bmt2
+                'PF04153' # NOT2/NOT3/NOT5 C-terminal
+        ]
+        
+
+        # prepare embeddings
+        #self.embedding_data = torch.load(args.swissprot_data_path)
+        # dataset iterator
+        #dataset = Facilitator_Dataset(args=args, dataset=self.embedding_data)
+        # create a clone of the dataset
+        #cloned_dataset = copy.deepcopy(dataset)
+
+        # Get indices and split them
+        #indices = list(range(len(dataset)))
+        #train_indices, valid_indices = train_test_split(indices, test_size=args.valid_size, random_state=args.seed)
+        
+        # create full dataloader
+        #self.all_dataloader = DataLoader(cloned_dataset, batch_size=args.batch_size, shuffle=False)
+        
+        # Create PyTorch DataLoader using the indices
+        #self.train_sampler = Subset(dataset, train_indices)
+        #self.valid_sampler = Subset(dataset, valid_indices)
+        #train_dataloader = DataLoader(train_sampler, batch_size=args.batch_size, shuffle=True)
+        #valid_dataloader = DataLoader(test_sampler, batch_size=args.batch_size, shuffle=False)
+    
+        ##########################################
+        # Load Stage 1 SwissProt+Pfam Embeddings #
+        ##########################################
+    
+        # initialize the embedding data to None
+        self.swissprot_data, self.pfam_data = None, None
+    
+        # get both the swissprot and pfam dataset iterator in one
+        if (args.swissprot_data_path != 'None') and (args.pfam_data_path != 'None'):
+            print('Load both SwissProt and Pfam dataset...')
+            self.train_dataset, self.valid_dataset, self.all_swiss_dataloader, self.all_pfam_dataloader = self.load_both()
+
+        # get the swissprot dataset iterator
+        elif args.pfam_data_path == 'None':
+            print('Load SwissProt dataset...')
+            self.train_dataset, self.valid_dataset, self.all_swiss_dataloader = self.load_swissprot()
+            self.all_pfam_dataloader = None
+
+        # get the pfam dataset iterator 
+        elif args.swissprot_data_path == 'None':
+            print('Load Pfam dataset...')
+            self.train_dataset, self.valid_dataset, self.all_pfam_dataloader = self.load_pfam()
+            self.all_swiss_dataloader = None
+            
+
+
+    def load_swissprot(self):
+
+        # prepare embeddings
+        self.swissprot_data = torch.load(self.args.swissprot_data_path)
+        
+        # dataset iterator
+        swiss_dataset = Facilitator_Dataset(args=self.args, dataset=self.swissprot_data)      
+        # create a clone of the dataset
+        cloned_swiss_dataset = copy.deepcopy(swiss_dataset)
+
+        # Get indices and split them
+        indices = list(range(len(swiss_dataset)))
+        train_indices, valid_indices = train_test_split(indices, test_size=self.args.valid_size, random_state=self.args.seed)
+        
+        # Create Pytorch iterator using the indices
+        swiss_train_subset = Subset(swiss_dataset, train_indices)
+        swiss_valid_subset = Subset(swiss_dataset, valid_indices)
+
+        # Create Pytorch dataloader on all samples
+        swiss_all_dataloader = DataLoader(cloned_swiss_dataset, batch_size=self.args.batch_size, shuffle=False)
+
+        
+        return (
+                swiss_train_subset,
+                swiss_valid_subset,
+                swiss_all_dataloader
+        )
+    
+    
+    def load_pfam(self):
+
+        # prepare embeddings
+        self.pfam_data = torch.load(self.args.pfam_data_path)
+        
+        # dataset iterator
+        pfam_dataset = Facilitator_Dataset(args=self.args, dataset=self.pfam_data)      
+        # create a clone of the dataset
+        cloned_pfam_dataset = copy.deepcopy(pfam_dataset)
+
+        # Get indices and split them
+        indices = list(range(len(pfam_dataset)))
+        train_indices, valid_indices = train_test_split(indices, test_size=self.args.valid_size, random_state=self.args.seed)
+        
+        # Create Pytorch Dataloader using the indices
+        pfam_train_subset = Subset(pfam_dataset, train_indices)
+        pfam_valid_subset = Subset(pfam_dataset, valid_indices)
+
+        # Create Pytorch dataloader on all samples
+        pfam_all_dataloader = DataLoader(cloned_pfam_dataset, batch_size=self.args.batch_size, shuffle=False)
+
+        return (
+                pfam_train_subset,
+                pfam_valid_subset,
+                pfam_all_dataloader
+        )
+    
+
+    def load_both(self):
+
+        # get swissprot
+        swissprot_train_subset, swissprot_valid_subset, swissprot_all_dataloader = self.load_swissprot()
+
+        # get pfam
+        pfam_train_subset, pfam_valid_subset, pfam_all_dataloader = self.load_pfam()
+    
+        # combined subsets 
+        combined_train_subset = ConcatDataset([swissprot_train_subset, pfam_train_subset])
+        combined_valid_subset = ConcatDataset([swissprot_valid_subset, pfam_valid_subset])
+
+        return (
+                combined_train_subset,
+                combined_valid_subset,
+                swissprot_all_dataloader,
+                pfam_all_dataloader
+        )
+
+
+    def train_dataloader(self):
+        return DataLoader(
+                self.train_dataset,
+                #self.train_sampler,
+                batch_size=self.args.batch_size,
+                #num_workers=self.args.num_workers,
+                shuffle=True,
+                #pin_memory=True
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+                self.valid_dataset,
+                #self.valid_sampler,
+                batch_size=self.args.batch_size,
+                #num_workers=self.args.num_workers,
+                #pin_memory=True
+        )
+
+    def test_dataloader(self):
+        # Define test dataloader if needed
+        pass
+
+

stage1_config.json

@@ -0,0 +1,50 @@
+{
+    "data_path": "None",
+    "pfam_data_path": "None",
+    "tb_logger_path": "None",
+    "tb_logger_folder": "None",
+    "version_name": "None",
+    "model_checkpoint_path": "/project/andrewferguson/niksapraljak/Project_ProtARDM/logs/Stage1_final_models/checkpoints/Pretraining_PENCiL_45M/epoch=19-step=116600.ckpt",
+    "output_dict_path": "/project/ranganathanr/niksapraljak/BioM3_PDZ/outputs/output_dict.pt",
+    "valid_size": 0.2,
+    "epochs": 10,
+    "acc_grad_batches": 1,
+    "batch_size": 80,
+    "num_workers": 12,
+    "weight_decay": "5e-7",
+    "patience": 1,
+    "factor": 0.8,
+    "temperature": 0.8,
+    "seed": 42,
+    "num_gpus": 1,
+    "precision": "16",
+    "dataset_type": "default",
+    "model_type": "pfam",
+    "fast_dev_run": 0,
+    "sequence_keyword": "protein_sequence",
+    "id_keyword": "primary_Accession",
+    "dataset_source": "swissprot",
+    "pfam_data_split_label": "0",
+    "base_lr": 0.0016,
+    "global_batch_size": 80,
+    "lr": 0.0005,
+    "seq_model_path": "/project/ranganathanr/niksapraljak/TextDiff_model_weights/Stage_1/pretrained_models/esm2_t33_650M_UR50D.pt",
+    "pretrained_seq": true,
+    "trainable_seq": true,
+    "rep_layer": 33,
+    "protein_encoder_embedding": 1280,
+    "protein_encoder_lr": 0.0005,
+    "pLM_n_layers_to_finetune": 1,
+    "text_model_path": "/project/ranganathanr/niksapraljak/TextDiff_model_weights/Stage_1/pretrained_models/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",
+    "pretrained_text": true,
+    "trainable_text": true,
+    "text_encoder_embedding": 768,
+    "text_encoder_lr": 0.0005,
+    "text_max_length": 512,
+    "bLM_n_layers_to_finetune": 1,
+    "proj_embedding_dim": 512,
+    "dropout": 0.1,
+    "head_lr": 0.0005,
+    "inference_data_path": "/project/ranganathanr/niksapraljak/BioM3_PDZ/data/test_prompts_PDZ_swissprot_pfam_dataset.csv",
+    "inference_output_path": "/project/ranganathanr/niksapraljak/BioM3_PDZ/outputs/Stage1_test_prompts_PDZ.pt"
+}