Upload 4 files

scripts/diffusion.py

@@ -0,0 +1,264 @@
+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

@@ -0,0 +1,129 @@
+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

@@ -0,0 +1,153 @@
+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

@@ -0,0 +1,165 @@
+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")