Blame for upstreams/duo/trainer_base.py

import itertools
from dataclasses import dataclass

import hydra.utils
import lightning as L
import numpy as np
import torch
import torch.nn.functional as F
import transformers

import dataloader
import metrics
import models
import utils


@dataclass
class Loss:
  loss: torch.FloatTensor
  nlls: torch.FloatTensor
  prior_loss: torch.FloatTensor
  num_tokens: torch.FloatTensor


class LogLinear(torch.nn.Module):
  def __init__(self, eps):
    super().__init__()
    self.eps = eps # 1e-3 by default, to be consistent with SEDD: https://github.com/louaaron/Score-Entropy-Discrete-Diffusion/blob/0605786da5ccb5747545e26d66fdf477187598b6/noise_lib.py#L56

  def forward(self, t):
    t = (1 - self.eps) * t
    alpha_t = 1 - t 
    dalpha_t = - torch.ones_like(alpha_t) * (1 - self.eps)
    return dalpha_t, alpha_t

  def get_t_for_alpha(self, alpha_t):
    return 1 - alpha_t


class Cosine(torch.nn.Module):
  def __init__(self, eps):
    super().__init__()
    self.eps = eps
    self.half_pi = torch.pi / 2

  def forward(self, t):
    t = (1 - self.eps) * t
    alpha_t = 1 - torch.cos(self.half_pi * (1 - t))
    dalpha_t = - torch.sin(self.half_pi * (1 - t)) * self.half_pi
    return dalpha_t, alpha_t

  def get_t_for_alpha(self, alpha_t):
    is_tensor = torch.is_tensor(alpha_t)
    if not is_tensor:
      alpha_t = torch.tensor([alpha_t])
    t = 1 - 2 / torch.pi * torch.acos(1 - alpha_t)
    if not is_tensor:
      t = t.cpu().item()
    return t


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 _unsqueeze(x, reference):
  return x.view(
    * x.shape,
    * ((1,) * (len(reference.shape) - len(x.shape))))


class TrainerBase(L.LightningModule):
  def __init__(
    self,
    config,
    tokenizer: transformers.PreTrainedTokenizer,
    vocab_size=None):
    super().__init__()
    self.save_hyperparameters()
    self.config = config
    if hasattr(self.config.algo, 'ignore_bos'):
      self.ignore_bos = config.algo.ignore_bos
    else:
      self.ignore_bos = False
    if hasattr(self.config.algo, 'loss_type'):
      self.loss_type = config.algo.loss_type
    self.tokenizer = tokenizer
    if vocab_size is None:
      self.vocab_size = len(self.tokenizer)
    else:
      self.vocab_size = vocab_size
    self.sampler = self.config.sampling.predictor
    self.antithetic_sampling = self.config.training.antithetic_sampling
    self.parameterization = self.config.algo.parameterization
    if self.config.algo.backbone == 'dit':
      self.backbone = models.dit.DIT(
        self.config, vocab_size=self.vocab_size)
    elif self.config.algo.backbone == 'unet':
      self.backbone = models.unet.UNet(self.config, self.vocab_size)
    elif self.config.algo.backbone == 'dimamba':
      self.backbone = models.dimamba.DiMamba(
        self.config,
        vocab_size=self.vocab_size,
        pad_token_id=self.tokenizer.pad_token_id)
    elif self.config.algo.backbone == 'hf_dit':
      self.backbone = transformers.AutoModelForMaskedLM.from_pretrained(
        config.eval.checkpoint_path, trust_remote_code=True)

    self.T = self.config.algo.T
    self.num_tokens = self.config.model.length
    self.p_nucleus = self.config.sampling.p_nucleus
    # Noise Schedule
    if config.noise.type == 'log-linear':
      self.noise = LogLinear(config.noise.eps)
    elif config.noise.type == 'cosine':
      self.noise = Cosine(config.noise.eps)
    else:
      raise ValueError(config.noise.type)
    # Class-conditional training arguments
    self.num_classes = config.data.get('num_classes', None)
    self.class_conditional = self.num_classes is not None
    self.class_cond_dropout = config.training.class_dropout_p 

    self.metrics = metrics.Metrics(
      gen_ppl_eval_model_name_or_path=\
        self.config.eval.gen_ppl_eval_model_name_or_path,
      eval_ppl_batch_size=\
        self.config.eval.perplexity_batch_size)

    if self.config.training.ema > 0:
      self.ema = models.ema.ExponentialMovingAverage(
        self._get_parameters(),
        decay=self.config.training.ema)
    else:
      self.ema = None
    
    self.lr = self.config.optim.lr
    self.sampling_eps = self.config.training.sampling_eps
    self.time_conditioning = self.config.algo.time_conditioning
    self.neg_infinity = -1000000.0
    self.fast_forward_epochs = None
    self.fast_forward_batches = None

  def _validate_configuration(self):
    assert self.config.algo.backbone in {'dit', 'hf_dit', 
                                         'unet'}
    if self.config.algo.parameterization == 'ar':
      assert not self.config.algo.time_conditioning
      assert self.config.prior.type == 'none'
    if self.time_conditioning:
      assert self.config.sampling.predictor != 'ancestral_cache'

    if self.parameterization in {'score', 'mean'}:
      assert self.time_conditioning
    if self.T > 0:
      assert self.parameterization != 'score'

  def to(self, *args, **kwargs):
    self = super().to(*args, **kwargs) 
    self.metrics.to(*args, **kwargs)
    return self

  def q_xt(self, x, alpha_t):
    raise NotImplementedError

  def _get_parameters(self):
    return itertools.chain(self.backbone.parameters(),
                           self.noise.parameters())

  def _eval_mode(self):
    if self.ema:
      self.ema.store(self._get_parameters())
      self.ema.copy_to(self._get_parameters())
    self.backbone.eval()
    self.noise.eval()

  def _train_mode(self):
    if self.ema:
      self.ema.restore(self._get_parameters())
    self.backbone.train()
    self.noise.train()

  def on_load_checkpoint(self, checkpoint):
    if self.ema:
      self.ema.load_state_dict(checkpoint['ema'])
    # Copied from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py#L41
    self.fast_forward_epochs = checkpoint['loops'][
      'fit_loop']['epoch_progress']['current']['completed']
    self.fast_forward_batches = checkpoint['loops'][
      'fit_loop']['epoch_loop.batch_progress'][
        'current']['completed']

  def on_save_checkpoint(self, checkpoint):
    if self.ema:
      checkpoint['ema'] = self.ema.state_dict()
    # Copied from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/tasks/seq.py
    # ['epoch_loop.batch_progress']['total']['completed']
    # is 1 iteration behind, so we're using the optimizer's progress.
    checkpoint['loops']['fit_loop'][
      'epoch_loop.batch_progress']['total'][
        'completed'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['total'][
              'completed'] * self.trainer.accumulate_grad_batches
    checkpoint['loops']['fit_loop'][
      'epoch_loop.batch_progress']['current'][
        'completed'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['current'][
              'completed'] * self.trainer.accumulate_grad_batches
    # _batches_that_stepped tracks the number of global steps,
    # not the number of local steps, so we don't multiply with
    # self.trainer.accumulate_grad_batches here.
    checkpoint['loops']['fit_loop'][
      'epoch_loop.state_dict'][
        '_batches_that_stepped'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['total']['completed']
    if 'sampler' not in checkpoint.keys():
      checkpoint['sampler'] = {}
    if hasattr(self.trainer.train_dataloader.sampler,
               'state_dict'):
      sampler_state_dict = self.trainer.\
        train_dataloader.sampler.state_dict()
      checkpoint['sampler'][
        'random_state'] = sampler_state_dict.get(
          'random_state', None)
    else:
      checkpoint['sampler']['random_state'] = None

  def on_train_start(self):
    if self.ema:
      self.ema.move_shadow_params_to_device(self.device)
    # Adapted from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py
    distributed = (
      self.trainer._accelerator_connector.use_distributed_sampler
      and self.trainer._accelerator_connector.is_distributed)
    if distributed:
      sampler_cls = dataloader.FaultTolerantDistributedSampler
    else:
      sampler_cls = dataloader.RandomFaultTolerantSampler
    updated_dls = []
    for dl in self.trainer.fit_loop._combined_loader.flattened:
      if hasattr(dl.sampler, 'shuffle'):
        dl_sampler = sampler_cls(
          dl.dataset, shuffle=dl.sampler.shuffle)
      else:
        dl_sampler = sampler_cls(dl.dataset)
      if (distributed
          and self.fast_forward_epochs is not None
          and self.fast_forward_batches is not None):
        dl_sampler.load_state_dict({
          'epoch': self.fast_forward_epochs,
          'counter': (self.fast_forward_batches
                      * self.config.loader.batch_size)})
      updated_dls.append(
        torch.utils.data.DataLoader(
          dl.dataset,
          batch_size=self.config.loader.batch_size,
          num_workers=self.config.loader.num_workers,
          pin_memory=self.config.loader.pin_memory,
          sampler=dl_sampler,
          shuffle=False,
          persistent_workers=True))
    self.trainer.fit_loop._combined_loader.flattened = updated_dls

  def optimizer_step(self, *args, **kwargs):
    super().optimizer_step(*args, **kwargs)
    if self.ema:
      self.ema.update(self._get_parameters())

  def _process_sigma(self, sigma):
    raise NotImplementedError

  def _process_model_output(self, model_output, xt, sigma):
    raise NotImplementedError

  def forward(self, xt, sigma, labels=None, weights=None,
              nn_input_idxs=None):
    if nn_input_idxs is None:
      nn_input_idxs = xt

    sigma = self._process_sigma(sigma)
    with torch.amp.autocast('cuda', dtype=torch.float32):
      model_output = self.backbone(
        x=nn_input_idxs, sigma=sigma, class_cond=labels, 
        weights=weights)
    return self._process_model_output(
      model_output=model_output, xt=xt, sigma=sigma)

  def on_train_epoch_start(self):
    self.metrics.reset()
    assert self.metrics.train_nlls.nll.mean_value == 0
    assert self.metrics.train_nlls.nll.weight == 0

  def training_step(self, batch, batch_idx):
    current_accumulation_step = (
      batch_idx % self.trainer.accumulate_grad_batches)
    losses = self._loss(batch['input_ids'],
                        batch.get('labels', None),
                        batch['attention_mask'],
                        current_accumulation_step,
                        train_mode=True)
    self.metrics.update_train(losses.nlls, losses.prior_loss,
                              losses.num_tokens)
    self.log(name='trainer/loss',
             value=losses.loss.item(),
             on_step=True,
             on_epoch=False,
             sync_dist=True)
    return losses.loss

  def on_train_epoch_end(self):
    for k, v in self.metrics.valid_nlls.items():
      self.log(name=k, value=v.compute(), on_step=False,
               on_epoch=True, sync_dist=True)

  def on_validation_epoch_start(self):
    self.metrics.reset()
    self._eval_mode()
    assert self.metrics.valid_nlls.nll.mean_value == 0
    assert self.metrics.valid_nlls.nll.weight == 0

  def validation_step(self, batch, batch_idx):
    del batch_idx
    losses = self._loss(batch['input_ids'],
                        batch.get('labels', None),
                        batch['attention_mask'])
    self.metrics.update_valid(losses.nlls, losses.prior_loss,
                              losses.num_tokens)
    return losses.loss

  def on_validation_epoch_end(self):
    for k, v in self.metrics.valid_nlls.items():
      self.log(name=k,  value=v.compute(), on_step=False,
               on_epoch=True, sync_dist=True)
    if ((self.config.eval.compute_perplexity_on_sanity
         or not self.trainer.sanity_checking)
         and self.config.eval.generate_samples):
      samples, text_samples = None, None
      for _ in range(
        self.config.sampling.num_sample_batches):
        samples = self.generate_samples(
          num_samples=self.config.loader.eval_batch_size)
        
        self.metrics.record_entropy(samples)
        # Decode the samples to be re-tokenized by eval model
        text_samples = self.tokenizer.batch_decode(samples)
        if self.config.eval.compute_generative_perplexity:
          self.metrics.record_generative_perplexity(
            text_samples, self.num_tokens, self.device)
      if text_samples is not None:
        if self.trainer.global_rank == 0 and hasattr(
          self.trainer.logger, 'log_table'):
          # Log the last generated samples
          text_samples = text_samples[
            : self.config.sampling.num_sample_log]
          self.trainer.logger.log_table(
            key=f'samples@global_step{self.global_step}',
            columns=['Generated Samples'],
            data=[[s] for s in text_samples])
        if self.config.eval.compute_generative_perplexity:
          self.log('val/gen_ppl',
                   self.metrics.gen_ppl.compute(),
                   on_epoch=True,
                   on_step=False,
                   sync_dist=True)
          self.log('val/sample_entropy',
                   self.metrics.sample_entropy.compute(),
                   on_epoch=True,
                   on_step=False,
                   sync_dist=True)
    self._train_mode()

  def configure_optimizers(self):
    optimizer = torch.optim.AdamW(
      self._get_parameters(),
      lr=self.config.optim.lr,
      betas=(self.config.optim.beta1,
             self.config.optim.beta2),
      eps=self.config.optim.eps,
      weight_decay=self.config.optim.weight_decay)

    scheduler = hydra.utils.instantiate(
      self.config.lr_scheduler, optimizer=optimizer)
    scheduler_dict = {'scheduler': scheduler,
                      'interval': 'step',
                      'monitor': 'val/loss',
                      'name': 'trainer/lr'}
    return [optimizer], [scheduler_dict]

  def generate_samples(self, num_samples, num_steps, eps):
    raise NotImplementedError

  def restore_model_and_sample(self, num_steps, eps=1e-5):
    """Generate samples from the model."""
    # Lightning auto-casting is not working in this method for some reason
    self._eval_mode()
    samples = self.generate_samples(
      num_samples=self.config.loader.eval_batch_size,
      num_steps=num_steps,
      eps=eps)
    self._train_mode()
    return samples

  def _process_model_input(self, x0, valid_tokens):
    raise NotImplementedError

  def nll(self, input_tokens, labels, output_tokens,
          current_accumulation_step=None, train_mode=False):
    raise NotImplementedError

  def _loss(self, x0, labels, valid_tokens,
            current_accumulation_step=None,
            train_mode=False):
    (input_tokens, output_tokens,
     valid_tokens) = self._process_model_input(
       x0, valid_tokens)
    loss = self.nll(input_tokens, labels, output_tokens,
                    current_accumulation_step, train_mode)
    assert loss.ndim == 2
    if self.ignore_bos:
      loss[:, 1:] = loss[:, 1:]
      valid_tokens[:, 1:] = valid_tokens[:, 1:]

    nlls = (loss * valid_tokens).sum()
    num_tokens = valid_tokens.sum()
    token_nll = nlls / num_tokens

    return Loss(loss=token_nll,
                nlls=nlls,
                prior_loss=0.0,
                num_tokens=num_tokens)


class Diffusion(TrainerBase):
  def _validate_configuration(self):
    super()._validate_configuration()
    assert self.config.sampling.noise_removal in {
      'none', 'ancestral', 'greedy'}
    assert self.loss_type in {'elbo', 'low_var'}
    if self.config.sampling.noise_removal == 'greedy':
      assert self.sampler != 'analytic'
      assert self.parameterization in {'mean', 'subs'}

  def _process_model_input(self, x0, valid_tokens):
    return x0, None, valid_tokens

  def _process_sigma(self, sigma):
    assert sigma.ndim == 2
    sigma = sigma.mean(-1).squeeze()
    if sigma.ndim == 0:
      sigma = sigma.unsqueeze(0)
    if not self.time_conditioning:
      sigma = torch.zeros_like(sigma)
    assert sigma.ndim == 1, sigma.shape
    return sigma

  def _sample_t(self, n, accum_step):
    if accum_step is not None:
      # During training
      batch_dim = n
      n = self.config.loader.global_batch_size
    _eps_t = torch.rand(n, device=self.device)
    if self.antithetic_sampling:
      offset = torch.arange(n, device=self.device) / n
      _eps_t = (_eps_t / n + offset) % 1
    t = (1 - self.sampling_eps) * _eps_t + self.sampling_eps
    if accum_step is not None:
      t = t.chunk(self.trainer.num_nodes)[self.trainer.node_rank]
      t = t.chunk(self.trainer.num_devices)[self.trainer.local_rank]
      t = t.chunk(self.trainer.accumulate_grad_batches)[
        accum_step]
      # corner case for the last datapoint
      t = t[:batch_dim]
    return t

  def _sigma_from_alphat(self, alpha_t):
    return -torch.log(alpha_t)

  def _reconstruction_loss(self, x0):
    t0 = torch.zeros(1, x0.shape[0], dtype=self.dtype,
                     device=self.device)
    sigma_t0 = self._sigma_from_alphat(self.noise(t0)[1])
    model_output_t0 = self.forward(x0, sigma_t0)
    return - torch.gather(input=model_output_t0,
                          dim=-1,
                          index=x0[:, :, None]).squeeze(-1)

  def nll_per_token(self, model_output, xt, x0, alpha_t,
                    dalpha_t, low_var):
    raise NotImplementedError

  def nll(self, x0, labels, output_tokens,
          current_accumulation_step=None, train_mode=False):
    del output_tokens
    t = self._sample_t(x0.shape[0],
                       current_accumulation_step)
    assert t.shape[0] == x0.shape[0]
    if self.T > 0:
      t = (t * self.T).to(torch.int)
      t = t / self.T
      # t \in {1/T, 2/T, ..., 1}
      t += (1 / self.T)
    
    dalpha_t, alpha_t = self.noise(t)
    alpha_t = alpha_t.unsqueeze(-1)
    dalpha_t = dalpha_t.unsqueeze(-1)
    assert alpha_t.ndim == 2
    assert dalpha_t.ndim == 2
    sigma = self._sigma_from_alphat(alpha_t)

    xt = self.q_xt(x0, alpha_t)
    # Handle class-conditional training, with class dropout
    if self.class_conditional:
      assert labels is not None
      rand = torch.rand(size=labels.shape, dtype=torch.float32, 
                      device=self.device)
      # num_classes represent the absence of class-conditioning
      labels = torch.where(rand < self.class_cond_dropout, 
                           self.num_classes, labels)
    else:
      assert labels is None
    log_x_theta = self.forward(xt, sigma=sigma, labels=labels)
    utils.print_nans(log_x_theta, 'model_output')
    return self.nll_per_token(
      log_x_theta=log_x_theta,
      xt=xt,
      x0=x0,
      alpha_t=alpha_t,
      dalpha_t=dalpha_t,
      low_var=train_mode and self.loss_type == 'low_var')

  def _get_score(self, **kwargs):
    del kwargs
    raise NotImplementedError

  def _denoiser_update(self, x, t):
    raise NotImplementedError

  def _analytic_update(self, x, t, dt):
    raise NotImplementedError

  def _posterior_from_x0(self, x0, xt, alpha_s, alpha_t):
    """From the clean x0, or denoiser predictions, implement 
    the mathematical expression for q(z_s | z_t, x)"""
    raise NotImplementedError

  def _forward_process(self, q_x0, alpha_s):
    """Apply forward noising: q(x_s | x_0), where x_0 is a 
       K-dimensional vector"""
    raise NotImplementedError

  def _get_ancestral_posterior(self, xt, sigma, labels, 
                              alpha_s, alpha_t, p_x0):
    """Top-level call from generate_samples. Compute the 
       standard posterior sampling distribution from the 
       noisy sequence xt"""
    gamma = self.config.sampling.guid_weight
    # If class cond but gamma is None or 0, same as sampling 
    #  from the class unconditional case.
    if not self.class_conditional or gamma in (None, 0.0):
      if labels is not None:
        labels = torch.full_like(labels, self.num_classes)
      return self._get_posterior_from_xt(xt, sigma, labels, 
                                         alpha_s, alpha_t, p_x0)
    elif gamma == 1.0:  
      # Simply sample from the cond. posterior only.
      assert labels is not None
      return self._get_posterior_from_xt(xt, sigma, labels, 
                                         alpha_s, alpha_t, p_x0)
    else:
      # Case gamma not in {0, 1}, and class conditional 
      #  -> mix conditional and unconditional predictions.
      return self._get_guided_posterior_from_xt(self, xt,
        sigma, labels, gamma, alpha_s, alpha_t, p_x0)

  def _get_posterior_from_xt(self, xt, sigma, labels, alpha_s, 
                             alpha_t, p_x0=None):
    """From xt, compute a single denoiser predictions, cast 
       to correct dtype, and call _posterior_from_x0."""
    if p_x0 is None:
      log_x0_pred = self.forward(xt, sigma, labels)

      if self.config.sampling.use_float64:
        log_x0_pred = log_x0_pred.to(torch.float64)

      if self.p_nucleus < 1:
        log_x0_pred = utils.top_k_top_p_filtering(
          log_x0_pred, top_p=self.p_nucleus)
      p_x0 = log_x0_pred.exp()

    return p_x0, self._posterior_from_x0(x0=p_x0,xt=xt,
      alpha_s=alpha_s, alpha_t=alpha_t)

  def _get_guided_posterior_from_xt(self, xt, sigma, labels, 
    gamma, alpha_s, alpha_t, p_x0=None):
    """From xt, combine the class cond / uncond predictions 
       of the denoiser. Call _get_posterior_from_xt."""
    # unpack the cache
    if p_x0 is None:
      p_x0_cond = p_x0_uncond = None
    else:
      p_x0_cond, p_x0_uncond = p_x0
    
    p_x0_cond, cond_posterior = self._get_posterior_from_xt(
      xt, sigma, labels, alpha_s, alpha_t, p_x0_cond)
    log_cond_posterior = cond_posterior.log()
    # NOTE: conditioning on self.num_classes represents the
    #  class-unconditional predictions.
    p_x0_uncond, uncond_posterior = self._get_posterior_from_xt(
      xt, sigma, torch.full_like(labels, self.num_classes), 
      alpha_s, alpha_t, p_x0_uncond)
    log_uncond_posterior = uncond_posterior.log()

    un_normalized_posterior = gamma * log_cond_posterior \
                       + (1 - gamma) * log_uncond_posterior
    # Handle cases where the posterior is zero (eg after 
    #  nucleus sampling)
    is_inf_mask = torch.logical_or(
      log_cond_posterior.isinf(), log_uncond_posterior.isinf())
    un_normalized_posterior[is_inf_mask] = self.neg_infinity
    
    return ((p_x0_cond, p_x0_uncond), 
            un_normalized_posterior.softmax(-1))

  def _ancestral_update(self, x, t, labels, dt, p_x0=None, 
                        noise_removal_step=False):
    _, alpha_t = self.noise(t)
    if noise_removal_step:
      alpha_s = torch.ones_like(alpha_t)
    else:
      _, alpha_s = self.noise(t - dt)
    assert alpha_t.ndim == 2

    sigma = self._sigma_from_alphat(alpha_t)
    assert alpha_t.ndim == 2, f'{alpha_t.ndim=}'

    p_x0, q_xs = self._get_ancestral_posterior(x, sigma, 
      labels, alpha_s, alpha_t, p_x0)

    return p_x0, sample_categorical(q_xs)

  def _psi_update(self, x, t, labels, dt, kappa, p_x0=None,
                  noise_removal_step=False):
    _, alpha_t = self.noise(t)
    if noise_removal_step:
      alpha_s = torch.ones_like(alpha_t)
    else:
      _, alpha_s = self.noise(t - dt)
    alpha_0 = torch.ones_like(alpha_t)
    sigma = self._sigma_from_alphat(alpha_t)

    # Standard posterior q(x_s | x_t)
    p_x0, q_xs = self._get_ancestral_posterior(
      x, sigma, labels, alpha_s, alpha_t, p_x0)

    # Posterior targeting t=0, reuse predictions p_x0
    _, q_x0 = self._get_ancestral_posterior(
      x, sigma, labels, alpha_0, alpha_t, p_x0)

    # PC: forward-noise q_x0 back to time s
    pc_q_xs = self._forward_process(q_x0, alpha_s)

    q_sample = kappa * q_xs + (1 - kappa) * pc_q_xs
    return p_x0, sample_categorical(q_sample)

  def _get_sampling_time_profile(self, eps, num_steps):
    profile = self.config.sampling.psi.time_profile
    num_steps += 1
    if profile == 'linear' \
      or self.config.sampling.predictor != 'psi':
      # Default: linearly decrease
      return torch.linspace(1, eps, num_steps)
    if not profile.startswith('linear-constant-linear'):
      raise ValueError(profile)
    c = float(profile.split('-')[3])
    if 'inv' in profile:
      c = self.noise.get_t_for_alpha(c)
    psi_cfg = self.config.sampling.psi
    n_hi = round(psi_cfg.high_frac * num_steps)
    n_mid = round(psi_cfg.middle_frac * num_steps)
    return torch.cat([
      torch.linspace(1, c, n_hi),
      torch.full((n_mid,), c),
      torch.linspace(c, eps, num_steps - n_hi - n_mid)])

  def _mode_to_psi_kappas(self, mode, timesteps):
    n = len(timesteps)
    if mode == 'pure-posterior':
      return torch.ones(n)
    if mode == 'pure-pc':
      return torch.zeros(n)

    eta = float(mode.split('-')[-1])
    if (mode.startswith('constant-')
        and not mode.startswith('constant-remdm')):
      return torch.full((n,), eta)

    # Noise-schedule-dependent modes (ReMDM-like)
    _, all_alphas = self.noise(timesteps)
    alpha_t, alpha_s = all_alphas[:-1], all_alphas[1:]
    eta_t = torch.tensor([eta])
    if mode.startswith('max-capped-'):
      sigma = torch.minimum(
        eta_t.expand_as(alpha_t), (1 - alpha_s) / alpha_t)
      sigma = torch.where(alpha_t == 0, eta, sigma)
    elif mode.startswith('max-rescale-'):
      sigma_max = torch.minimum(
        eta_t.expand_as(alpha_t), (1 - alpha_s) / alpha_t)
      sigma = torch.where(alpha_t > 0, sigma_max, 1) * eta
    elif mode.startswith('constant-remdm'):
      sigma = eta_t
    else:
      raise ValueError(mode)
    kappas = torch.clip(1 - sigma / (1 - alpha_s), 0, 1)
    if len(kappas) > 0:
      kappas = torch.cat([kappas, torch.ones(1)])
    return kappas

  def _get_kappas(self, timesteps):
    cfg = self.config.sampling.psi
    n = len(timesteps)
    n_hi = round(cfg.high_frac * n)
    n_mid = round(cfg.middle_frac * n)
    kappas = torch.cat([
      self._mode_to_psi_kappas(cfg.high_mode, timesteps[:n_hi]),
      self._mode_to_psi_kappas(cfg.middle_mode,
        timesteps[n_hi:n_hi + n_mid]),
      self._mode_to_psi_kappas(cfg.low_mode,
        timesteps[n_hi + n_mid:])])
    assert (kappas >= 0).all() and (kappas <= 1).all()
    assert len(kappas) == n
    return kappas

  @torch.no_grad()
  def generate_samples(self, num_samples, labels=None,
                       num_steps=None, eps=1e-5):
    """Generate samples from the model."""
    # Lightning auto-casting is not working in this method for some reason
    if num_steps is None:
      num_steps = self.config.sampling.steps
    x = self.prior_sample(num_samples, self.num_tokens)
    use_psi_sampler = self.config.sampling.predictor == 'psi'
    timesteps = self._get_sampling_time_profile(eps, 
                                                num_steps)
    if use_psi_sampler:
      kappas = self._get_kappas(timesteps).to(self.device)

    p_x0_cache = None
    if labels is not None:
      labels = labels.to(self.device)

    for i in range(num_steps):
      t = timesteps[i] * torch.ones(
        x.shape[0], 1, device=self.device)
      dt = timesteps[i] - timesteps[i+1]
      if self.sampler == 'ancestral':
        _, x = self._ancestral_update(
          x=x, t=t, labels=labels, dt=dt, p_x0=None)
      elif self.sampler == 'ancestral_cache':
        p_x0_cache, x_next = self._ancestral_update(
          x=x, t=t, labels=labels, dt=dt, p_x0=p_x0_cache)
        if (not torch.allclose(x_next, x)
            or self.time_conditioning):
          # Disable caching
          p_x0_cache = None
        x = x_next
      elif self.sampler == 'psi':
        _, x = self._psi_update(x=x, t=t, kappa=kappas[i],
                                labels=labels, dt=dt, 
                                p_x0=None)
      elif self.sampler == 'analytic':
        assert labels is None, 'class-conditional sampling ' \
          'is not implemented with the analytic sampler'
        x = self._analytic_update(x=x,t=t, dt=dt)
      else:
        raise ValueError(self.sampler)

    t0 = timesteps[-1] * torch.ones(x.shape[0], 1,
                                    device=self.device)
    if self.config.sampling.noise_removal == 'ancestral':
      if self.sampler == 'analytic':
        x = self._denoiser_update(x=x, t=t0)
      else:
        _, x = self._ancestral_update(x=x, t=t0, labels=labels,
          dt=None, p_x0=p_x0_cache, noise_removal_step=True)
    elif self.config.sampling.noise_removal == 'greedy':
      sigma = self._sigma_from_alphat(self.noise(t0)[1])
      x = self.forward(xt=x, sigma=sigma).argmax(dim=-1)
    return x

  @torch.no_grad
  def _semi_ar_sampler(
    self, n_samples, stride_length, num_strides, dt=0.001):
    # TODO(subham): Test this method after refactoring.
    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.prior_sample(n_samples, self.num_tokens)
      if target is not None:
        x[:, : -stride_length] = target
      for i in range(num_steps + 1):
        p_x0_cache, x_next = self._ancestral_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.batch_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
    # TODO(subham): Test this method after refactoring.
    self._eval_mode()
    (sampling_steps, samples,
     sequence_lengths) = self._semi_ar_sampler(
      n_samples=self.config.loader.eval_batch_size,
      stride_length=stride_length,
      num_strides=num_strides, 
      dt=dt)
    self._train_mode()
    return sampling_steps, samples, sequence_lengths


class AbsorbingState(Diffusion):
  def __init__(self, config, tokenizer):
    # NOTE: Ideally, we should do 
    # vocab_size = len(tokenizer), so that we account
    # for the special tokens added in dataloader.py.
    # But we use tokenizer.vocab_size so as to to be
    # consistent with the prior checkpoints.
    vocab_size = tokenizer.vocab_size
    if (not hasattr(tokenizer, 'mask_token')
        or tokenizer.mask_token is None):
      self.mask_index = vocab_size
      vocab_size += 1
    else:
      self.mask_index = tokenizer.mask_token_id
    self.subs_masking = config.algo.subs_masking
    super().__init__(config, tokenizer,
                     vocab_size=vocab_size)
    self.save_hyperparameters()

  def _validate_configuration(self):
    super()._validate_configuration()
    if self.parameterization in {'score', 'mean'}:
      assert self.time_conditioning
    assert not (self.parameterization == 'mean'
                and self.T == 0)
    if self.T > 0:
      assert self.parameterization in {'mean', 'subs'}
    if self.subs_masking:
      assert self.parameterization == 'mean'

  def q_xt(self, x, alpha_t):
    """Computes the noisy sample xt.

    Args:
      x: int torch.Tensor with shape (batch_size,
          diffusion_model_input_length), input. 
      alpha_t: float torch.Tensor with shape (batch_size, 1).
    """
    move_indices = torch.rand(
      * x.shape, device=x.device) < 1 - alpha_t
    xt = torch.where(move_indices, self.mask_index, x)
    if self.ignore_bos:
      xt[:, 0] = x[:, 0]
    return xt

  def prior_sample(self, *batch_dims):
    return self.mask_index * torch.ones(
      * batch_dims, dtype=torch.int64, device=self.device)

  def _posterior_from_x0(self, x0, xt, alpha_s, alpha_t):
    """From the clean x0, or denoiser predictions, implement 
    the mathematical expression for q(z_s | z_t, x)"""
    assert x0.dtype == torch.float64, 'Requires float64 prec.'
    # should be one-hot on clean tokens
    orig_mask = xt[:, :, None] != self.mask_index
    orig_mask = orig_mask.expand(-1, -1, x0.shape[-1])
    orig_output_on_clean = x0[orig_mask]

    q_xs = ((alpha_s - alpha_t) / (1 - alpha_t))[..., None] * x0
    q_xs[..., self.mask_index] = (1 - alpha_s) / (1 - alpha_t)
    q_xs[orig_mask] = orig_output_on_clean
    return q_xs

  def _forward_process(self, x0, alpha_s):
    out = alpha_s[..., None] * x0
    out[..., self.mask_index] = 1 - alpha_s
    return out

  def _staggered_score(self, score, dsigma):
    score = score.clone()
    extra_const = (1 - dsigma.exp()) * score.sum(dim=-1)
    score *= dsigma.exp()[:, None]
    score[..., self.mask_index] += extra_const
    return score

  def _analytic_update(self, x, t, dt):
    sigma_t = self._sigma_from_alphat(self.noise(t)[1])
    sigma_s = self._sigma_from_alphat(self.noise(t - dt)[1])
    dsigma = sigma_t - sigma_s
    score = self._get_score(x, sigma_t)
    if self.config.sampling.use_float64:
      score = score.to(torch.float64)
    stag_score = self._staggered_score(score, dsigma)
    probs = stag_score * self._transp_transition(x, dsigma)
    return sample_categorical(probs)

  def _denoiser_update(self, x, t):
    sigma = self._sigma_from_alphat(self.noise(t)[1])
    score = self._get_score(x, sigma)
    if self.config.sampling.use_float64:
      score = score.to(torch.float64)
    stag_score = self._staggered_score(score, sigma)
    probs = stag_score * self._transp_transition(x, sigma)
    probs[..., self.mask_index] = 0
    samples = sample_categorical(probs)
    return samples

  def _transp_transition(self, i, sigma):
    sigma = _unsqueeze(sigma, reference=i[..., None])
    edge = torch.exp(-sigma) * F.one_hot(
      i, num_classes=self.vocab_size)
    edge += torch.where(i == self.mask_index,
                        1 - torch.exp(-sigma).squeeze(-1),
                        0)[..., None]
    return edge


class UniformState(Diffusion):
  def _validate_configuration(self):
    super()._validate_configuration()
    assert self.time_conditioning
    assert self.parameterization == 'mean'
    if self.config.algo.name != 'distillation':
      assert self.T == 0

  def _forward_process(self, x0, alpha_s):
    return (alpha_s[..., None] * x0
            + (1 - alpha_s[..., None]) / self.vocab_size)

  def q_xt(self, x, alpha_t):
    """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) < 1 - alpha_t
    uniform_tensor = torch.randint(
      0, self.vocab_size, x.shape, device=x.device)
    xt = torch.where(move_indices, uniform_tensor, x)
    if self.ignore_bos:
      xt[:, 0] = x[:, 0]
    return xt

  def prior_sample(self, *batch_dims):
    return torch.randint(
      0, self.vocab_size, batch_dims, dtype=torch.int64,
      device=self.device)