Delete scripts

scripts/diffusion.py

@@ -1,264 +0,0 @@
-import itertools
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-import pytorch_lightning as L
-import torchmetrics
-from dataclasses import dataclass
-from esm_utils import load_esm2_model
-from transformers import AutoModelForMaskedLM, AutoModel, AutoTokenizer
-import dit, ema
-import sys
-import config
-import wandb
-import noise_schedule  # Assuming this is part of the MDLM repository
-
-wandb_key = "2b76a2fa2c1cdfddc5f443602c17b011fefb0a8f"
-wandb.login(key=wandb_key)
-wandb.init(project=config.Wandb.PROJECT, group=config.Wandb.GROUP)
-
-LOG2 = math.log(2)
-
-# Goal is to build an MDLM head using pre-existing ESM LM head
-# Wrap the ESM model to obtain logits and ignore sigma to work with MDLM codebase
-class WrapESM(nn.Module):
-    def __init__(self, esm_model_path=config.MODEL_NAME):
-        super(WrapESM, self).__init__()
-        self.model = AutoModelForMaskedLM.from_pretrained(esm_model_path)
-        self.tokenizer = AutoTokenizer.from_pretrained(esm_model_path)
-    
-    # def __getattr__(self, name):
-    #     return getattr(self.model, name)
-
-    def __call__(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
-    
-    def freeze_model(self):
-       # Disable parameter updates for all layers
-        for param in self.model.parameters():
-            param.requires_grad = False
-
-    def unfreeze_n_layers(self):
-        # Count number of encoder layers
-        model_layers = len(self.model.esm.encoder.layer)
-
-        # Enable parameter updates for the last 3 encoder layers
-        for i, layer in enumerate(self.model.esm.encoder.layer):
-            if i >= model_layers-config.ESM_LAYERS:
-                for module in layer.attention.self.key.modules():
-                    for param in module.parameters():
-                        param.requires_grad = True
-                for module in layer.attention.self.query.modules():
-                    for param in module.parameters():
-                        param.requires_grad = True
-                for module in layer.attention.self.value.modules():
-                    for param in module.parameters():
-                        param.requires_grad = True
-        
-    def forward(self, sigma, **inputs):
-        return self.model(**inputs)
-
-    def save_model(self, save_dir):
-        self.model.save_pretrained(save_dir)
-        self.tokenizer.save_pretrained(save_dir)
-
-    def load_model(self, load_dir):
-        self.model = AutoModel.from_pretrained(load_dir)
-        self.tokenizer = AutoTokenizer.from_pretrained(load_dir)
-
-@dataclass
-class Loss:
-    loss: torch.FloatTensor
-    nlls: torch.FloatTensor
-    token_mask: torch.FloatTensor
-
-class NLL(torchmetrics.MeanMetric):
-    pass
-
-class BPD(NLL):
-    def compute(self) -> torch.Tensor:
-        """Computes the bits per dimension.
-        Returns:
-          bpd
-        """
-        return self.mean_value / self.weight / LOG2
-
-class Perplexity(NLL):
-    def compute(self) -> torch.Tensor:
-        """Computes the Perplexity.
-        Returns:
-         Perplexity
-        """
-        return torch.exp(self.mean_value / self.weight)
-
-
-# Based on MDLM repo
-class Diffusion(L.LightningModule):
-    def __init__(self, config, tokenizer):
-        super().__init__()
-        self.config = config
-        self.tokenizer = tokenizer
-
-        self.softplus = torch.nn.Softplus()
-        metrics = torchmetrics.MetricCollection({
-            'nll': NLL(),
-            'bpd': BPD(),
-            'ppl': Perplexity(),
-        })
-        metrics.set_dtype(torch.float64)
-        self.train_metrics = metrics.clone(prefix='train/')
-        self.valid_metrics = metrics.clone(prefix='val/')
-        self.test_metrics = metrics.clone(prefix='test/')
-
-        self.T = self.config.T
-        self.lr = self.config.Optim.LR
-        self.backbone = WrapESM(self.config.MODEL_NAME)
-        self.noise = noise_schedule.get_noise(self.config, dtype=self.dtype)
-        self.time_conditioning = self.config.TIME_CONDITIONING
-        self.subs_masking = self.config.SUBS_MASKING
-        self.mask_index = self.tokenizer.mask_token_id
-        self.antithetic_sampling = self.config.Training.ANTITHETIC_SAMPLING
-        self.sampling_eps = self.config.Training.SAMPLING_EPS
-        self.neg_infinity = -1000000.0
-
-
-    ############ FORWARD DIFFUSION #########
-    def compute_loss(self, latents, attention_mask, val):
-        """"Average of MLM losses to stabilize training"""
-        self.noise.eval() if val else self.noise.train()
-        loss = self.forward_diffusion(latents)
-
-        nlls = loss * attention_mask
-        count = attention_mask.sum()
-        batch_nll = nlls.sum()
-        token_nll = batch_nll / count
-
-        return Loss(loss=token_nll, nlls=nlls, token_mask=attention_mask)
-
-    def forward_diffusion(self, x0):
-        """Forward diffusion process, adds noise to the latents."""
-        t = self.sample_timestep(x0.shape[0], x0.device)
-
-        sigma, dsigma = self.noise(t)
-        unet_conditioning = sigma[:, None]
-        move_chance = 1 - torch.exp(-sigma[:, None])
-
-        xt = self.q_xt(x0, move_chance)
-        model_output = self.forward(xt, unet_conditioning)
-
-        # SUBS parameterization, continuous time.
-        log_p_theta = torch.gather(input=model_output, dim=-1, index=x0[:, :, None]).squeeze(-1)
-        scale = (dsigma / torch.expm1(sigma))[:, None]
-        loss = - log_p_theta * scale
-        return loss
-
-    def sample_timestep(self, n, device):
-        _eps_t = torch.rand(n, device=device)
-        if self.antithetic_sampling:
-            offset = torch.arange(n, device=device) / n
-            _eps_t = (_eps_t / n + offset) % 1
-        t = (1 - self.sampling_eps) * _eps_t + self.sampling_eps
-        return t
-    
-    def q_xt(self, x, move_chance):
-        """
-        Computes the noisy sample xt.
-        Args:
-            x: int torch.Tensor with shape (batch_size, diffusion_model_input_length), input. 
-            move_chance: float torch.Tensor with shape (batch_size, 1).
-        """
-        move_indices = torch.rand(* x.shape, device=x.device) < move_chance
-        xt = torch.where(move_indices, self.mask_index, x) # Use variable masking rate to mask tokens (introduce noise)
-        return xt
-
-    def forward(self, latents, sigma):
-        esm_outputs = self.backbone(latents, sigma)
-        optimized_logits = self.subs_parameterization(esm_outputs.logits, latents)
-        return optimized_logits
-
-    def subs_parameterization(self, logits, xt):
-        logits[:, :, self.mask_index] += self.neg_infinity
-        logits = logits - torch.logsumexp(logits, dim=-1, keepdim=True)
-
-        unmasked_indices = (xt != self.mask_index)
-        logits[unmasked_indices] = self.neg_infinity
-        logits[unmasked_indices, xt[unmasked_indices]] = 0
-        return logits
-
-
-    ######### GENERATION #########
-    def sample_prior(self, *batch_dims):
-        return self.mask_index * torch.ones(* batch_dims, dtype=torch.int64)
-
-    def sample_categorical(categorical_probs):
-        gumbel_norm = (1e-10 - (torch.rand_like(categorical_probs) + 1e-10).log())
-        return (categorical_probs / gumbel_norm).argmax(dim=-1)
-
-    def ddpm_caching_update(self, x, t, dt, p_x0=None):
-        sigma_t, _ = self.noise(t)
-        if t.ndim > 1:
-            t = t.squeeze(-1)
-        assert t.ndim == 1
-        move_chance_t = t[:, None, None]
-        move_chance_s = (t - dt)[:, None, None]
-        assert move_chance_t.ndim == 3, move_chance_t.shape
-        if p_x0 is None:
-            p_x0 = self.forward(x, sigma_t).exp()
-    
-        assert move_chance_t.ndim == p_x0.ndim
-        q_xs = p_x0 * (move_chance_t - move_chance_s)
-        q_xs[:, :, self.mask_index] = move_chance_s[:, :, 0]
-        _x = self.sample_categorical(q_xs)
-        
-        copy_flag = (x != self.mask_index).to(x.dtype)
-        return p_x0, copy_flag * x + (1 - copy_flag) * _x
-
-
-    @torch.no_grad()
-    def sample_subs_guidance(self, n_samples, stride_length, num_strides, dt=0.001):
-        ones = torch.ones(n_samples, dtype=self.dtype,device=self.device)
-        num_steps = int(1 / dt)
-        sampling_steps = 0
-        intermediate_tokens = []
-        target = None
-
-        for _ in range(num_strides + 1):
-            p_x0_cache = None
-            x = self.sample_prior(n_samples,self.config.model.length).to(self.device)
-            
-            if target is not None:
-                x[:, : -stride_length] = target
-            
-            for i in range(num_steps + 1):
-                p_x0_cache, x_next = self.ddpm_caching_update(x=x, t=(1 - i * dt) * ones, dt=dt, p_x0=p_x0_cache)
-                if (not torch.allclose(x_next, x) or self.time_conditioning):
-                    p_x0_cache = None
-                    sampling_steps += 1
-                x = x_next
-            x = self.forward(x, 0 * ones).argmax(dim=-1)
-            intermediate_tokens.append(x[:, :stride_length].cpu().numpy())
-            target = x[:, stride_length:]
-    
-        intermediate_tokens.append(target.cpu().numpy())
-        intermediate_text_samples = []
-        sequence_lengths = ((np.concatenate(intermediate_tokens, axis=1)[:, 1:]
-                                 == self.tokenizer.eos_token_id).cumsum(-1) == 0).sum(-1)
-    
-        for i in range(2, len(intermediate_tokens) + 1):
-            intermediate_text_samples.append(self.tokenizer.decode(np.concatenate(intermediate_tokens[:i], axis=1)))
-        
-        return (sampling_steps, intermediate_text_samples,
-            sequence_lengths)
-
-    def restore_model_and_semi_ar_sample(self, stride_length, num_strides, dt=0.001):
-        """Generate samples from the model."""
-        # Lightning auto-casting is not working in this method for some reason
-        self.backbone.eval()
-        self.noise.eval()
-    
-        (sampling_steps, samples, sequence_lengths) = self.sample_subs_guidance(n_samples=self.config.Loader.BATCH_SIZE,stride_length=stride_length,num_strides=num_strides,dt=dt)
-
-        self.backbone.train()
-        self.noise.train()
-        return sampling_steps, samples, sequence_lengths

scripts/generate.py

@@ -1,129 +0,0 @@
-import torch
-import numpy as np
-from transformers import AutoTokenizer, AutoModel
-from diffusion import Diffusion
-import config
-from esm_utils import load_esm2_model, get_latents
-
-def mask_sequence(sequence, mask_char='X'):
-    """Masks parts of the sequence based on the mask_char."""
-    mask_indices = [i for i, char in enumerate(sequence) if char == mask_char]
-    masked_sequence = sequence.replace(mask_char, '[MASK]')
-    return masked_sequence, mask_indices
-
-def generate_filled_sequence(model, tokenizer, esm_model, masked_sequence, mask_indices):
-    """Generates the filled sequence for the masked regions."""
-    inputs = tokenizer(masked_sequence, return_tensors="pt")
-    with torch.no_grad():
-        outputs = esm_model(**inputs)
-    latents = outputs.last_hidden_state.squeeze(0)
-    
-    sigma = torch.rand(1, device=latents.device)
-    noisy_latents = model.forward(latents, sigma)
-    denoised_latents = model.reverse_diffusion(noisy_latents, sigma)
-    
-    filled_sequence = list(masked_sequence)
-    for idx in mask_indices:
-        token_id = torch.argmax(denoised_latents[idx]).item()
-        filled_sequence[idx] = tokenizer.decode([token_id])
-    
-    return ''.join(filled_sequence)
-
-def generate_scaffold_sequence(model, tokenizer, esm_model, peptides, final_length):
-    """Generates a scaffold sequence to connect multiple peptides."""
-    total_peptide_length = sum(len(peptide) for peptide in peptides)
-    scaffold_length = final_length - total_peptide_length
-    if scaffold_length <= 0:
-        raise ValueError("Final length must be greater than the combined length of the peptides.")
-    
-    scaffold = "[MASK]" * scaffold_length
-    masked_sequence = "".join(peptides[:1] + [scaffold] + peptides[1:])
-    
-    inputs = tokenizer(masked_sequence, return_tensors="pt")
-    with torch.no_grad():
-        outputs = esm_model(**inputs)
-    latents = outputs.last_hidden_state.squeeze(0)
-    
-    sigma = torch.rand(1, device=latents.device)
-    noisy_latents = model.forward(latents, sigma)
-    denoised_latents = model.reverse_diffusion(noisy_latents, sigma)
-    
-    filled_sequence = list(masked_sequence)
-    scaffold_start = len(peptides[0])
-    scaffold_end = scaffold_start + scaffold_length
-    for idx in range(scaffold_start, scaffold_end):
-        token_id = torch.argmax(denoised_latents[idx]).item()
-        filled_sequence[idx] = tokenizer.decode([token_id])
-    
-    return ''.join(filled_sequence)
-
-def generate_de_novo_sequence(model, tokenizer, esm_model, sequence_length):
-    """Generates a de novo protein sequence of the specified length."""
-    scaffold = "[MASK]" * sequence_length
-    masked_sequence = scaffold
-    
-    inputs = tokenizer(masked_sequence, return_tensors="pt")
-    with torch.no_grad():
-        outputs = esm_model(**inputs)
-    latents = outputs.last_hidden_state.squeeze(0)
-    
-    sigma = torch.rand(1, device=latents.device)
-    noisy_latents = model.forward(latents, sigma)
-    denoised_latents = model.reverse_diffusion(noisy_latents, sigma)
-    
-    filled_sequence = list(masked_sequence)
-    for idx in range(sequence_length):
-        token_id = torch.argmax(denoised_latents[idx]).item()
-        filled_sequence[idx] = tokenizer.decode([token_id])
-    
-    return ''.join(filled_sequence)
-
-if __name__ == "__main__":
-    import argparse
-
-    # Argument parsing
-    parser = argparse.ArgumentParser(description="Generate protein sequences using latent diffusion model.")
-    subparsers = parser.add_subparsers(dest="mode")
-
-    # Subparser for the first strategy (multiple peptides to scaffold)
-    parser_scaffold = subparsers.add_parser("scaffold", help="Generate scaffold to connect multiple peptides.")
-    parser_scaffold.add_argument("peptides", nargs='+', help="Peptides to connect.")
-    parser_scaffold.add_argument("final_length", type=int, help="Final length of the protein sequence.")
-
-    # Subparser for the second strategy (fill in regions)
-    parser_fill = subparsers.add_parser("fill", help="Fill in specified regions in a given protein sequence.")
-    parser_fill.add_argument("sequence", help="Protein sequence with regions to fill specified by 'X'.")
-
-    # Subparser for the third strategy (de novo generation)
-    parser_de_novo = subparsers.add_parser("de_novo", help="Generate a de novo protein sequence.")
-    parser_de_novo.add_argument("sequence_length", type=int, help="Length of the de novo generated protein sequence.")
-
-    args = parser.parse_args()
-
-    # Load models
-    tokenizer, esm_model = load_esm2_model(config.MODEL_NAME)
-    diffusion_model = Diffusion.load_from_checkpoint(config.Training.SAVE_DIR + "best_model.ckpt", config=config, latent_dim=config.LATENT_DIM)
-    diffusion_model.eval()
-
-    if args.mode == "scaffold":
-        peptides = args.peptides
-        final_length = args.final_length
-        filled_sequence = generate_scaffold_sequence(diffusion_model, tokenizer, esm_model, peptides, final_length)
-        print(f"Peptides:          {' '.join(peptides)}")
-        print(f"Final Length:      {final_length}")
-        print(f"Generated Protein: {filled_sequence}")
-
-    elif args.mode == "fill":
-        sequence = args.sequence
-        masked_sequence, mask_indices = mask_sequence(sequence)
-        filled_sequence = generate_filled_sequence(diffusion_model, tokenizer, esm_model, masked_sequence, mask_indices)
-        print(f"Original Sequence: {sequence}")
-        print(f"Masked Sequence:   {masked_sequence}")
-        print(f"Filled Sequence:   {filled_sequence}")
-
-    elif args.mode == "de_novo":
-        sequence_length = args.sequence_length
-        filled_sequence = generate_de_novo_sequence(diffusion_model, tokenizer, esm_model, sequence_length)
-        print(f"De Novo Sequence Length: {sequence_length}")
-        print(f"Generated Protein: {filled_sequence}")
-

scripts/noise_schedule.py

@@ -1,153 +0,0 @@
-import abc
-
-import torch
-import torch.nn as nn
-
-# Flags required to enable jit fusion kernels
-torch._C._jit_set_profiling_mode(False)
-torch._C._jit_set_profiling_executor(False)
-torch._C._jit_override_can_fuse_on_cpu(True)
-torch._C._jit_override_can_fuse_on_gpu(True)
-
-
-def get_noise(config, dtype=torch.float32):
-  return LogLinearNoise()
-
-  if config.noise.type == 'geometric':
-    return GeometricNoise(config.noise.sigma_min,
-                          config.noise.sigma_max)
-  elif config.noise.type == 'loglinear':
-    return LogLinearNoise()
-  elif config.noise.type == 'cosine':
-    return CosineNoise()
-  elif config.noise.type == 'cosinesqr':
-    return CosineSqrNoise()
-  elif config.noise.type == 'linear':
-    return Linear(config.noise.sigma_min,
-                  config.noise.sigma_max,
-                  dtype)
-  else:
-    raise ValueError(f'{config.noise.type} is not a valid noise')
-
-
-def binary_discretization(z):
-  z_hard = torch.sign(z)
-  z_soft = z / torch.norm(z, dim=-1, keepdim=True)
-  return z_soft + (z_hard - z_soft).detach()
-
-
-class Noise(abc.ABC, nn.Module):
-  """
-  Baseline forward method to get the total + rate of noise at a timestep
-  """
-  def forward(self, t):
-    # Assume time goes from 0 to 1
-    return self.total_noise(t), self.rate_noise(t)
-  
-  @abc.abstractmethod
-  def rate_noise(self, t):
-    """
-    Rate of change of noise ie g(t)
-    """
-    pass
-
-  @abc.abstractmethod
-  def total_noise(self, t):
-    """
-    Total noise ie \int_0^t g(t) dt + g(0)
-    """
-    pass
-
-
-class CosineNoise(Noise):
-  def __init__(self, eps=1e-3):
-    super().__init__()
-    self.eps = eps
-
-  def rate_noise(self, t):
-    cos = (1 - self.eps) * torch.cos(t * torch.pi / 2)
-    sin = (1 - self.eps) * torch.sin(t * torch.pi / 2)
-    scale = torch.pi / 2
-    return scale * sin / (cos + self.eps)
-
-  def total_noise(self, t):
-    cos = torch.cos(t * torch.pi / 2)
-    return - torch.log(self.eps + (1 - self.eps) * cos)
-
-
-class CosineSqrNoise(Noise):
-  def __init__(self, eps=1e-3):
-    super().__init__()
-    self.eps = eps
-
-  def rate_noise(self, t):
-    cos = (1 - self.eps) * (
-      torch.cos(t * torch.pi / 2) ** 2)
-    sin = (1 - self.eps) * torch.sin(t * torch.pi)
-    scale = torch.pi / 2
-    return scale * sin / (cos + self.eps)
-
-  def total_noise(self, t):
-    cos = torch.cos(t * torch.pi / 2) ** 2
-    return - torch.log(self.eps + (1 - self.eps) * cos)
-
-
-class Linear(Noise):
-  def __init__(self, sigma_min=0, sigma_max=10, dtype=torch.float32):
-    super().__init__()
-    self.sigma_min = torch.tensor(sigma_min, dtype=dtype)
-    self.sigma_max = torch.tensor(sigma_max, dtype=dtype)
-
-  def rate_noise(self, t):
-    return self.sigma_max - self.sigma_min
-
-  def total_noise(self, t):
-    return self.sigma_min + t * (self.sigma_max - self.sigma_min)
-
-  def importance_sampling_transformation(self, t):
-    f_T = torch.log1p(- torch.exp(- self.sigma_max))
-    f_0 = torch.log1p(- torch.exp(- self.sigma_min))
-    sigma_t = - torch.log1p(- torch.exp(t * f_T + (1 - t) * f_0))
-    return (sigma_t - self.sigma_min) / (
-      self.sigma_max - self.sigma_min)
-
-
-class GeometricNoise(Noise):
-  def __init__(self, sigma_min=1e-3, sigma_max=1):
-    super().__init__()
-    self.sigmas = 1.0 * torch.tensor([sigma_min, sigma_max])
-
-  def rate_noise(self, t):
-    return self.sigmas[0] ** (1 - t) * self.sigmas[1] ** t * (
-      self.sigmas[1].log() - self.sigmas[0].log())
-
-  def total_noise(self, t):
-    return self.sigmas[0] ** (1 - t) * self.sigmas[1] ** t
-
-
-class LogLinearNoise(Noise):
-  """Log Linear noise schedule.
-  
-  Built such that 1 - 1/e^(n(t)) interpolates between 0 and
-  ~1 when t varies from 0 to 1. Total noise is
-  -log(1 - (1 - eps) * t), so the sigma will be
-  (1 - eps) * t.
-  """
-  def __init__(self, eps=1e-3):
-    super().__init__()
-    self.eps = eps
-    self.sigma_max = self.total_noise(torch.tensor(1.0))
-    self.sigma_min = self.eps + self.total_noise(torch.tensor(0.0))
-
-  def rate_noise(self, t):
-    return (1 - self.eps) / (1 - (1 - self.eps) * t)
-
-  def total_noise(self, t):
-    return -torch.log1p(-(1 - self.eps) * t)
-
-  def importance_sampling_transformation(self, t):
-    f_T = torch.log1p(- torch.exp(- self.sigma_max))
-    f_0 = torch.log1p(- torch.exp(- self.sigma_min))
-    sigma_t = - torch.log1p(- torch.exp(t * f_T + (1 - t) * f_0))
-    t = - torch.expm1(- sigma_t) / (1 - self.eps)
-    return t

scripts/train_pytorch.py

@@ -1,165 +0,0 @@
-import torch
-import config
-import math
-import sys
-import os
-from tqdm import tqdm
-from torch.optim import AdamW
-from transformers import AutoTokenizer
-from diffusion import WrapESM, Diffusion
-from data_loader import get_dataloaders
-
-def save_hyperparams(ckpt_dir):
-    hyperparms_txt_file = os.path.join(ckpt_dir, "hyperparameters.txt")
-    with open(hyperparms_txt_file, 'w') as f:
-        for k, v in vars(config).items():
-            if k.isupper():
-                f.write(f"{k}: {v}\n")
-
-def train_and_validate(model, optimizer, device, train_loader, val_loader, num_epochs, ckpt_dir):
-    best_val_loss = float('inf')
-
-    for epoch in range(num_epochs):
-        model.train()
-
-        print(f"EPOCH {epoch+1}/{num_epochs}")
-        sys.stderr.flush()
-        total_loss = 0.0
-        train_tokens = 0
-        weighted_total_train_loss = 0.0
-
-        train_update_interval = len(train_loader) // 4
-
-        with tqdm(enumerate(train_loader), desc="Training batch", total=len(train_loader), leave=True, position=0, ncols=100) as trainbar:
-            for step, inputs in trainbar:
-                inputs = {k: v.to(device) for k, v in inputs.items()}
-                optimizer.zero_grad()
-                outputs = model(**inputs)
-                train_loss = diffusion_model.compute_loss(inputs["input_ids"], inputs['attention_mask'],
-                                                          val=False).loss
-                train_loss.backward()
-                optimizer.step()
-
-                total_loss += train_loss.item()
-                weighted_total_train_loss += train_loss.item() * inputs['input_ids'].shape[1] # Loss * sequence length
-                train_tokens += inputs['input_ids'].shape[1]
-
-                if (step+1) % train_update_interval == 0:
-                    trainbar.update(train_update_interval)
-
-            avg_train_loss = total_loss / len(train_loader)
-            avg_train_neg_log_likelihood = weighted_total_train_loss / train_tokens
-            train_perplexity = math.exp(avg_train_neg_log_likelihood)
-
-        # Save model every epoch
-        train_ckpt_path = os.path.join(config.Eval.CHECKPOINT_PATH, f'epoch{epoch+1}')
-        model.save_model(train_ckpt_path)
-        save_hyperparams(train_ckpt_path)
-
-        # Validate model
-        if val_loader:
-            model.eval()
-            total_val_loss = 0.0
-            weighted_total_val_loss = 0.0
-            val_tokens = 0
-
-            with torch.no_grad():
-                val_update_interval = len(val_loader) // 4
-
-                with tqdm(enumerate(val_loader), desc='Validiation batch', total=len(val_loader), leave=True, position=0) as valbar:
-                    for step, inputs in valbar:
-                        inputs = {k: v.to(device) for k, v in inputs.items()}
-                        outputs = model(**inputs)
-                        val_loss = diffusion_model.compute_loss(inputs['input_ids'], inputs['attention_mask'],
-                                                                val=True).loss.item()
-                        
-                        total_val_loss += val_loss
-                        weighted_total_val_loss += val_loss * inputs['input_ids'].shape[1] # Loss * sequence length
-                        val_tokens += inputs['input_ids'].shape[1]
-
-                        if (step+1) % val_update_interval == 0:
-                            valbar.update(val_update_interval)
-
-                avg_val_loss = total_val_loss / len(val_loader)
-                avg_val_log_likelihood = weighted_total_val_loss / val_tokens
-                val_perplexity = math.exp(avg_val_log_likelihood)
-
-        # Save the best model based on validation loss
-        if avg_val_loss < best_val_loss:
-            best_val_loss = avg_val_loss
-            val_ckpt_path = os.path.join(config.Eval.CHECKPOINT_PATH, "best_model_epoch")
-            model.save_model(val_ckpt_path)
-            save_hyperparams(val_ckpt_path)
-
-
-        print(f"Average train loss: {avg_train_loss}")
-        print(f"Average train perplexity: {train_perplexity}\n")
-        sys.stdout.flush()
-
-        print(f"Average validation loss: {avg_val_loss}")
-        print(f"Average validation perplexity: {val_perplexity}\n")
-        sys.stdout.flush()
-        
-
-    return avg_train_loss, train_perplexity, avg_val_loss, val_perplexity
-                            
-
-def test(model, test_loader, device):
-    model.to(device).eval()
-    total_test_loss = 0.0
-    weighted_total_test_loss = 0.0
-    test_tokens = 0
-
-    with torch.no_grad():
-        for step, inputs in enumerate(test_loader):
-            inputs = {k: v.to(device) for k, v in inputs.items()}
-            outputs = model(**inputs)
-            test_loss = diffusion_model.compute_loss(inputs['input_ids'], inputs['attention_mask'],
-                                                     val=True).loss.item()
-
-            total_test_loss += test_loss
-            weighted_total_test_loss += test_loss * inputs['input_ids'].shape[1] # loss * sequence length
-            test_tokens += inputs['input_ids'].shape[1]
-        
-        avg_test_loss = total_test_loss / len(test_loader)
-        avg_test_log_likelihood = weighted_total_test_loss / test_tokens
-        test_perplexity = math.exp(avg_test_log_likelihood)
-
-    return avg_test_loss, test_perplexity
-
-
-if __name__ == "__main__":
-    device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
-    tokenizer = AutoTokenizer.from_pretrained(config.MODEL_NAME)
-
-    esm_model = WrapESM()
-    diffusion_model = Diffusion(config, tokenizer=tokenizer)
-
-    print(f'Trainable params before unfreezing: {sum(p.numel() for p in esm_model.parameters() if p.requires_grad)}')
-    
-    esm_model.to(device)
-    diffusion_model.to(device)
-
-    esm_model.freeze_model()
-    esm_model.unfreeze_n_layers()
-
-    print(f'Trainable params after unfreezing: {sum(p.numel() for p in esm_model.parameters() if p.requires_grad)}')
-
-    train_loader, val_loader, test_loader = get_dataloaders(config)
-    optimizer = AdamW(filter(lambda p: p.requires_grad, esm_model.parameters()), lr=config.Optim.LR)
-    
-    # Train and test the model
-    avg_train_loss, train_ppl, avg_val_loss, val_ppl = train_and_validate(esm_model, optimizer, device, train_loader, val_loader, config.Training.NUM_EPOCHS, config.Eval.CHECKPOINT_PATH)
-    avg_test_loss, test_ppl = test(esm_model, test_loader, device)
-
-    results_dict = {"Average train loss": avg_train_loss,
-                    "Average train perplexity": train_ppl,
-                    "Average val loss": avg_val_loss,
-                    "Average val perplexity": val_ppl,
-                    "Average test loss": avg_test_loss,
-                    "Average test perplexity": test_ppl,
-    }
-
-    print("TRAIN AND TEST RESULTS")
-    for k, v in results_dict.items():
-        print(f"{k}: {v}\n")