import models
import os
import time
import pdb
import collections
import copy
import pickle
import fsspec
import numpy as np
import torch
import torch.nn.functional as F
from torch.distributions import Normal
import trainer_base
import utils
import torch
from omegaconf import DictConfig
# torch.serialization.add_safe_globals([DictConfig])
class AR(trainer_base.TrainerBase):
def __init__(self, config, tokenizer):
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
super().__init__(config, tokenizer, vocab_size=vocab_size)
self.save_hyperparameters()
self._validate_configuration()
self.temp = config.algo.get("temp", 1.0)
def _validate_configuration(self):
super()._validate_configuration()
assert not self.config.algo.time_conditioning
assert self.config.prior.type == "none"
def _process_model_input(self, x0, valid_tokens):
input_tokens = x0[:, :-1]
output_tokens = x0[:, 1:]
valid_tokens = valid_tokens[:, 1:]
return input_tokens, output_tokens, valid_tokens
def nll(self, input_tokens, output_tokens, current_accumulation_step, *args):
del current_accumulation_step
output = self.backbone(input_tokens, None)
output[:, :, self.mask_index] = self.neg_infinity
output = output.log_softmax(-1)
return -output.gather(-1, output_tokens[:, :, None])[:, :, 0]
def generate_samples(self, num_samples, prompt_mask=None, prompt_tokens=None, **kwargs):
# precompute token buffer
num_pred_tokens = self.num_tokens - 1
x = torch.zeros(
(num_samples, num_pred_tokens + 1), dtype=torch.long, device=self.device
)
if prompt_mask is not None and prompt_tokens is not None:
x = (prompt_mask * prompt_tokens + (1 - prompt_mask) * x).long()
start = prompt_mask.sum(1).max().long().item() - 1
else:
start = 0
x[:, 0] = self.tokenizer.bos_token_id
# precompute noise
noise = (
torch.distributions.Gumbel(0, 1)
.sample((num_samples, num_pred_tokens, self.vocab_size))
.to(self.device)
)
if self.config.sampling.use_float64:
noise = noise.to(torch.float64)
for i in range(start, num_pred_tokens):
output = self.backbone(x[:, : i + 1], None)
output[:, :, self.mask_index] = self.neg_infinity
output = output / self.temp
output = output.log_softmax(-1)
y = (output[:, -1, :] + noise[:, i, :]).argmax(-1)
x[:, i + 1] = y
return x
def _process_sigma(self, sigma):
del sigma
return None
class MDLM(trainer_base.AbsorbingState):
def __init__(self, config, tokenizer):
super().__init__(config, tokenizer)
self._validate_configuration()
def _validate_configuration(self):
# ancestral sampling isn't desirable because it's slow
assert self.sampler == "ancestral_cache" or self.sampler == "posterior"
def _process_model_output(self, model_output, xt, sigma):
del sigma
model_output = model_output / self.temp
model_output[:, :, self.mask_index] += self.neg_infinity
# Normalize the model_output such that x.exp() is
# a probability distribution over vocab_size.
model_output = model_output - torch.logsumexp(
model_output, dim=-1, keepdim=True
)
# Apply updates directly in the logits matrix.
# For the logits of the unmasked tokens, set all values
# to -infinity except for the indices corresponding to
# the unmasked tokens.
unmasked_indices = xt != self.mask_index
model_output[unmasked_indices] = self.neg_infinity
model_output[unmasked_indices, xt[unmasked_indices]] = 0
return model_output
def nll_per_token(self, log_x_theta, xt, x0, alpha_t, dalpha_t, low_var=False):
del xt
log_p_theta = torch.gather(
input=log_x_theta, dim=-1, index=x0[:, :, None]
).squeeze(-1)
return log_p_theta * dalpha_t / (1 - alpha_t)
def _get_score(self, x, sigma):
model_output = self.forward(x, sigma)
# score(x, t) = p_t(y) / p_t(x)
# => log score(x, t) = log p_t(y) - log p_t(x)
# case 1: x = masked
# (i) y = unmasked
# log score(x, t) = log p_\theta(x)|_y + log k
# where k = exp(- sigma) / (1 - exp(- sigma))
# (ii) y = masked
# log score(x, t) = 0
# case 2: x = unmasked
# (i) y != masked, y != x
# log score(x_i, t) = - inf
# (ii) y = x
# log score(x_i, t) = 0
# (iii) y = masked token
# log score(x_i, t) = - log k
# where k = exp(- sigma) / (1 - exp(- sigma))
log_k = -torch.log(torch.expm1(sigma)).squeeze(-1)
assert log_k.ndim == 1
masked_score = model_output + log_k[:, None, None]
masked_score[:, :, self.mask_index] = 0
unmasked_score = self.neg_infinity * torch.ones_like(model_output)
unmasked_score = torch.scatter(
unmasked_score, -1, x[..., None], torch.zeros_like(unmasked_score[..., :1])
)
unmasked_score[:, :, self.mask_index] = -(log_k[:, None] * torch.ones_like(x))
masked_indices = (x == self.mask_index).to(model_output.dtype)[:, :, None]
model_output = masked_score * masked_indices + unmasked_score * (
1 - masked_indices
)
return model_output.exp()
class D3PMAbsorb(trainer_base.AbsorbingState):
def __init__(self, config, tokenizer):
super().__init__(config, tokenizer)
self._validate_configuration()
def _validate_configuration(self):
super()._validate_configuration()
assert self.noise.type == "log-linear"
assert self.parameterization == "mean"
def _process_model_output(self, model_output, xt, sigma):
del xt
del sigma
if self.subs_masking:
model_output[:, :, self.mask_index] += self.neg_infinity
return model_output.log_softmax(dim=-1)
def nll_per_token(self, log_x_theta, xt, x0, alpha_t, dalpha_t, low_var=False):
del dalpha_t
assert not low_var
dt = 1 / self.T
t = 1 - alpha_t # Only valid for log-linear schedule.
t = t.clamp(0.0, 1.0 - 1e-4)
alpha_t = alpha_t + torch.zeros_like(xt)
alpha_s = t - dt + torch.zeros_like(xt)
assert alpha_s.shape == xt.shape
assert alpha_t.shape == xt.shape
log_x_theta_at_x0 = torch.gather(log_x_theta, -1, x0[:, :, None]).squeeze(-1)
log_x_theta_at_m = log_x_theta[:, :, self.mask_index]
x_theta_at_m = log_x_theta_at_m.exp()
term_1_coef = dt / t
term_1_log_nr = torch.log(alpha_t * x_theta_at_m / t + 1)
term_1_log_dr = log_x_theta_at_x0
term_2_coef = 1 - dt / t
term_2_log_nr = term_1_log_nr
term_2_log_dr = torch.log(alpha_s * x_theta_at_m / (t - dt) + 1)
L_vb_masked = term_1_coef * (term_1_log_nr - term_1_log_dr) + term_2_coef * (
term_2_log_nr - term_2_log_dr
)
diffusion_loss = self.T * L_vb_masked * (xt == self.mask_index)
return self._reconstruction_loss(x0) + diffusion_loss
class SEDDAbsorb(trainer_base.AbsorbingState):
def __init__(self, config, tokenizer):
super().__init__(config, tokenizer)
self._validate_configuration()
def _validate_configuration(self):
super()._validate_configuration()
assert self.config.sampling.predictor == "analytic"
def _get_score(self, x, sigma):
return self.forward(x, sigma).exp()
def _process_model_output(self, model_output, xt, sigma):
esigm1_log = (
torch.where(sigma < 0.5, torch.expm1(sigma), sigma.exp() - 1)
.log()
.to(model_output.dtype)
)
# logits shape
# (batch_size, context_length, vocab_size)
model_output = (
model_output
- esigm1_log[:, None, None]
- np.log(model_output.shape[-1] - 1)
)
# The below scatter operation sets the log score
# for the input word to 0.
model_output = torch.scatter(
model_output, -1, xt[..., None], torch.zeros_like(model_output[..., :1])
)
return model_output
def nll_per_token(self, log_x_theta, xt, x0, alpha_t, dalpha_t, low_var=False):
"""Computes the SEDD loss for the Absorbing State Diffusion.
Args:
log_x_theta: float torch.Tensor with shape (batch_size,
context_length, vocab_size),
log score, output of the denoising network.
xt: int torch.Tensor with shape (batch_size,
context_length), input.
x0: int torch.Tensor with shape (batch_size,
context_length), input.
alpha_t: float torch.Tensor with shape (batch_size, 1),
signal level.
alpha_t: float torch.Tensor with shape (batch_size, 1),
signal level.
dalpha_t: float or float torch.Tensor with shape (batch_size, 1),
time derivative of signal level.
low_var: bool, low variance loss during training.
Returns:
loss with shape (batch_size, context_length).
"""
assert not low_var
masked_indices = xt == self.mask_index
sigma = self._sigma_from_alphat(alpha_t)
dsigma = -dalpha_t / alpha_t
expsig_minus_1 = torch.expm1(sigma).expand_as(xt)
q_ratio = 1 / expsig_minus_1[masked_indices]
words_that_were_masked = x0[masked_indices]
neg_term = q_ratio * torch.gather(
log_x_theta[masked_indices], -1, words_that_were_masked[..., None]
).squeeze(-1)
score = log_x_theta[masked_indices].exp()
if self.mask_index == self.vocab_size - 1:
pos_term = score[:, :-1].sum(dim=-1)
else:
pos_term = score[:, : self.mask_index].sum(dim=-1) + score[
:, self.mask_index + 1 :
].sum(dim=-1)
const = q_ratio * (q_ratio.log() - 1)
entropy = torch.zeros(*xt.shape, device=xt.device)
entropy[masked_indices] += pos_term - neg_term + const
return dsigma * entropy
class DUO_BASE(trainer_base.UniformState):
def __init__(self, config, tokenizer):
self.temp = config.algo.get("temp", 1.0)
super().__init__(config, tokenizer)
self._validate_configuration()
def on_save_checkpoint(self, checkpoint):
checkpoint["state_dict"] = collections.OrderedDict(
(k, v)
for k, v in checkpoint["state_dict"].items()
if not k.startswith("teacher")
)
super().on_save_checkpoint(checkpoint)
def on_load_checkpoint(self, checkpoint):
checkpoint["state_dict"] = collections.OrderedDict(
(k, v)
for k, v in checkpoint["state_dict"].items()
if not k.startswith("teacher")
)
super().on_load_checkpoint(checkpoint)
def _process_model_output(self, model_output, xt, sigma):
del xt, sigma
return (model_output / self.temp).log_softmax(dim=-1)
def _compute_posterior(self, x, xt, alpha_s, alpha_t):
"""Computes the posterior / approximate posterior.
Args:
x: Either clean input `x0` (one-hot),
or model's predicted `x_theta` of shape (B, L, V).
xt: The noisy latent (as indices) of shape (B, L).
alpha_s: Noise level at s of shape (B, [L | 1], 1).
alpha_t: Noise level at t of shape (B, [L | 1], 1).
Returns:
Posterior / approximate posterior of shape (B, L, V).
"""
if self.config.sampling.use_float64:
x = x.to(torch.float64)
if alpha_s.ndim == 2:
alpha_s = alpha_s.unsqueeze(-1)
if alpha_t.ndim == 2:
alpha_t = alpha_t.unsqueeze(-1)
alpha_ts = alpha_t / alpha_s
d_alpha = alpha_s - alpha_t
xt_one_hot = F.one_hot(xt, self.vocab_size).to(self.dtype).to(self.device)
return (
alpha_t * self.vocab_size * x * xt_one_hot
+ (alpha_ts - alpha_t) * xt_one_hot
+ d_alpha * x
+ (1 - alpha_ts) * (1 - alpha_s) / self.vocab_size
) / (
alpha_t * self.vocab_size * torch.gather(x, -1, xt[..., None])
+ (1 - alpha_t)
)
def nll_per_token(self, log_x_theta, xt, x0, alpha_t, dalpha_t, low_var=False):
assert alpha_t.ndim == 2
assert x0.ndim == 2
assert xt.ndim == 2
assert not torch.is_tensor(dalpha_t) or dalpha_t.ndim == 2
x_reconst = log_x_theta.exp()
x_bar_theta = (
self.vocab_size * alpha_t[:, :, None] * x_reconst + 1 - alpha_t[:, :, None]
)
coeff = dalpha_t / (self.vocab_size * alpha_t)
x_eq_xt = (x0 == xt).float()
x_neq_xt = 1 - x_eq_xt
xbar_xt = (1 - alpha_t) + self.vocab_size * alpha_t * x_eq_xt
xbar_theta_xt = torch.gather(x_bar_theta, -1, xt.unsqueeze(-1)).squeeze(-1)
xbar_theta_x = torch.gather(x_bar_theta, -1, x0.unsqueeze(-1)).squeeze(-1)
term1 = self.vocab_size * (1 / xbar_xt - 1 / xbar_theta_xt)
const = (1 - alpha_t) / (self.vocab_size * alpha_t + 1 - alpha_t)
term2_coefs = x_eq_xt * const + x_neq_xt
term2_offset = (
(self.vocab_size - 1) * const * x_eq_xt - (1 / const) * x_neq_xt
) * const.log()
term2_theta = -term2_coefs * (
x_bar_theta.log().sum(-1) - self.vocab_size * xbar_theta_xt.log()
)
term2_theta = (
term2_theta
- self.vocab_size
* alpha_t
/ (1 - alpha_t)
* (xbar_theta_x.log() - xbar_theta_xt.log())
* x_neq_xt
)
term2 = term2_theta + term2_offset
diffusion_loss = coeff * (term1 - term2)
assert diffusion_loss.ndim == 2
return diffusion_loss
def _ancestral_update(self, x, t, dt, p_x0=None, noise_removal_step=False, prompt_mask=None):
del p_x0
_, alpha_t = self.noise(t)
if noise_removal_step:
alpha_s = torch.ones_like(alpha_t)
else:
_, alpha_s = self.noise(t - dt)
sigma_t = self._sigma_from_alphat(alpha_t)
assert alpha_t.ndim == 2
cur_logits = (self.forward(x, sigma_t)).double()
xtheta_to_use = cur_logits.exp()
q_xs = self._compute_posterior(
x=xtheta_to_use, xt=x, alpha_s=alpha_s, alpha_t=alpha_t
)
if self.p_nucleus < 1:
q_xs = utils.top_k_top_p_filtering(q_xs.log(), top_p=self.p_nucleus)
xs = trainer_base.sample_categorical(q_xs)
# to handle conditional generation for inputs / infilling tasks
if prompt_mask is not None:
# expects the original x to have the prompt tokens where prompt_mask = 1
xs = (xs * (1 - prompt_mask) + x * prompt_mask).long()
return None, xs
class Integral(torch.autograd.Function):
"""
torch module calculating UDLM's p_t
"""
@staticmethod
def forward(ctx, gamma_t, data):
gamma_max = data["gamma_max"]
gamma_min = data["gamma_min"]
if (gamma_t.max() > gamma_max) or (gamma_t.min() < gamma_min):
print("max:{} {}".format(gamma_t.max(), gamma_max))
print("min:{} {}".format(gamma_t.min(), gamma_min))
gamma_t = torch.clip(gamma_t, gamma_min, gamma_max)
indices = torch.round(
(data["num_points"] - 1) * (gamma_t - gamma_min) / (gamma_max - gamma_min)
).long()
grad_pt = data["grad_pt"]
ctx.grad_pt = grad_pt[indices]
pt = data["pt"][indices]
assert pt.shape == gamma_t.shape
return pt
@staticmethod
def backward(ctx, grad_output):
return ctx.grad_pt * grad_output, None
class DUO(DUO_BASE):
def __init__(self, config, tokenizer):
super().__init__(config, tokenizer)
with fsspec.open(self.config.algo.integral_cache_path, "rb") as f:
self.integral_cache = pickle.load(f)
self.integral_cache["pt"] = torch.from_numpy(self.integral_cache["pt"])
self.integral_cache["grad_pt"] = torch.from_numpy(
self.integral_cache["grad_pt"]
)
self.gamma_min = self.config.algo.gamma_min
self.gamma_max = self.config.algo.gamma_max
self.gumbel_tau_log10_start = self.config.algo.gumbel_tau_log10_start
self.gumbel_tau_log10_end = self.config.algo.gumbel_tau_log10_end
self.curriculum_start = self.config.algo.curriculum_start
self.curriculum_end = self.config.algo.curriculum_end
self.loss_type = self.config.algo.loss_type
self.temp = config.algo.temp
self._validate_configuration()
def to(self, *args, **kwargs):
self = super().to(*args, **kwargs)
self.integral_cache["pt"] = self.integral_cache["pt"].to(*args, **kwargs)
self.integral_cache["grad_pt"] = self.integral_cache["grad_pt"].to(
*args, **kwargs
)
return self
def _compute_gumbel_tau_inverse(self):
start = self.gumbel_tau_log10_start
end = self.gumbel_tau_log10_end
delta = end - start
if self.global_step < self.curriculum_start:
tau = start
elif self.global_step < self.curriculum_end:
frac = (self.global_step - self.curriculum_start) / (
self.curriculum_end - self.curriculum_start
)
tau = start + frac * delta
else:
tau = -10
return 10 ** (-tau)
def training_step(self, batch, batch_idx):
self.log(
name="gumbel_tau_log10",
value=1 / self._compute_gumbel_tau_inverse(),
on_step=True,
on_epoch=False,
sync_dist=True,
)
return super().training_step(batch, batch_idx)
def _gamma_to_alphat(self, gamma_t):
integral = Integral.apply(gamma_t, self.integral_cache)
return (self.vocab_size * integral - 1) / (self.vocab_size - 1)
def _prior_loss(self):
alpha_1 = self._gamma_to_alphat(torch.tensor(self.gamma_max))
loss = (alpha_1 + (1 - alpha_1) / self.vocab_size) * torch.log(
(self.vocab_size - 1) * alpha_1 + 1
) + (1 - 1 / self.vocab_size) * (1 - alpha_1) * torch.log(1 - alpha_1)
return loss.item()
def _q_xt_gaussian(self, x, gamma_t):
"""Computes the noisy sample xt."""
assert gamma_t.ndim == 1
assert x.ndim == 3
gamma_t = gamma_t.unsqueeze(-1).unsqueeze(-1)
alpha_t = torch.sigmoid(-gamma_t).sqrt()
sigma_t = torch.sigmoid(gamma_t).sqrt()
epsilon = torch.randn(x.shape, dtype=torch.float32, device=self.device)
return alpha_t * x + sigma_t * epsilon
def nll(self, x0, output_tokens, current_accumulation_step=None, train_mode=False):
use_true_nll = self.global_step > self.curriculum_end or not train_mode
if use_true_nll:
return super().nll(x0, output_tokens, current_accumulation_step)
del output_tokens
t = self._sample_t(x0.shape[0], current_accumulation_step)
gamma_t = self.gamma_min + t * (self.gamma_max - self.gamma_min)
gamma_t_prime = self.gamma_max - self.gamma_min
usdm_alpha_t = self._gamma_to_alphat(gamma_t)
T = 1000
usdm_dalpha_t = (
gamma_t_prime * T * (self._gamma_to_alphat(gamma_t + 1 / T) - usdm_alpha_t)
)
usdm_alpha_t = usdm_alpha_t.unsqueeze(-1)
usdm_dalpha_t = usdm_dalpha_t.unsqueeze(-1)
assert usdm_alpha_t.ndim == 2
sigma = self._sigma_from_alphat(usdm_alpha_t)
x0_one_hot = F.one_hot(x0, self.vocab_size)
xt = self._q_xt_gaussian(x0_one_hot, gamma_t)
xt = xt * self._compute_gumbel_tau_inverse()
xt_usdm = xt.argmax(-1)
log_x_theta = self.forward(xt, sigma=sigma)
return self.nll_per_token(
log_x_theta=log_x_theta,
xt=xt_usdm,
x0=x0,
alpha_t=usdm_alpha_t,
dalpha_t=usdm_dalpha_t,
low_var=False,
)
class Distillation(DUO):
def __init__(self, config, tokenizer):
super().__init__(config, tokenizer)
self.update_teacher_every = config.algo.update_teacher_every
self.save_hyperparameters()
self.teacher = None
self.teacher_ema = config.algo.teacher_ema
self.linear_growth_dt = config.algo.linear_growth_dt
self.linear_growth_min = config.algo.linear_growth_min
self.linear_growth_max = config.algo.linear_growth_max
def _validate_configuration(self):
assert os.path.exists(self.config.algo.integral_cache_path), (
"The integral cache (Eq. 10 in the paper) for "
f"the {self.config.data.tokenizer_name_or_path} "
" tokenizer doesnt exist at "
f"{self.config.algo.integral_cache_path}. "
"Please generate it by running the utils.py script, "
"and ensure the correct path is specified using the "
"algo.integral_cache_path flag."
)
assert self.loss_type in {"kl-fwd", "kl-bwd", "posterior", "kl-posterior"}
def _maybe_update_teacher_weights(self):
if self.global_step % self.update_teacher_every != 0:
return
if self.teacher_ema:
self.ema.copy_to(self.teacher.parameters())
else:
for better_param, current_param in zip(
self.backbone.parameters(), self.teacher.parameters()
):
if current_param.requires_grad:
current_param.data.copy_(better_param.data)
@torch.no_grad()
def _teacher_logits(self, xt, sigma):
if self.teacher is None:
self.teacher = copy.deepcopy(self.backbone)
self._maybe_update_teacher_weights()
sigma = self._process_sigma(sigma)
with torch.cuda.amp.autocast(dtype=torch.float32):
model_output = self.teacher(xt, sigma)
logits = self._process_model_output(
model_output=model_output, xt=xt, sigma=sigma
)
return logits.detach()
def _sample_trajectory(self, x0, gamma_t, gamma_s):
"""Computes the noisy sample xt."""
assert gamma_t.ndim == 1
assert gamma_s.ndim == 1
assert x0.ndim == 2
x0 = F.one_hot(x0, self.vocab_size).to(self.dtype).to(self.device)
gamma_t = gamma_t.unsqueeze(-1).unsqueeze(-1)
alpha_t = torch.sigmoid(-gamma_t).sqrt()
sigma_t = torch.sigmoid(gamma_t).sqrt()
gamma_s = gamma_s.unsqueeze(-1).unsqueeze(-1)
alpha_s = torch.sigmoid(-gamma_s).sqrt()
sigma_s = torch.sigmoid(gamma_s).sqrt()
epsilon = torch.randn(x0.shape, dtype=torch.float32, device=self.device)
xt = alpha_t * x0 + sigma_t * epsilon
xs = alpha_s * x0 + sigma_s * epsilon
return xt, xs
def _compute_dt(self):
if self.linear_growth_dt:
scale = self.global_step / self.trainer.max_steps
return self.linear_growth_min + scale * (
self.linear_growth_max - self.linear_growth_min
)
n = self.global_step // self.update_teacher_every
return 2**n / self.T
def nll(self, x0, output_tokens, current_accumulation_step=None, train_mode=None):
del output_tokens, train_mode
t = self._sample_t(x0.shape[0], current_accumulation_step)
dt = self._compute_dt()
t = torch.clip(t + dt, 0, 1)
gamma_t = self.gamma_min + t * (self.gamma_max - self.gamma_min)
gamma_s = self.gamma_min + (t - dt) * (self.gamma_max - self.gamma_min)
usdm_alpha_t = self._gamma_to_alphat(gamma_t)
usdm_alpha_t = usdm_alpha_t.unsqueeze(-1)
assert usdm_alpha_t.ndim == 2
usdm_alpha_s = self._gamma_to_alphat(gamma_s)
usdm_alpha_s = usdm_alpha_s.unsqueeze(-1)
assert usdm_alpha_s.ndim == 2
xt, xs = self._sample_trajectory(x0, gamma_t, gamma_s)
xt_discrete = xt.argmax(-1)
xs_discrete = xs.argmax(-1)
log_x_theta_student = self.forward(
xt_discrete, sigma=self._sigma_from_alphat(usdm_alpha_t)
)
log_x_theta_teacher = self._teacher_logits(
xs_discrete, sigma=self._sigma_from_alphat(usdm_alpha_s)
)
if self.config.training.loss_precision == "float64":
log_x_theta_student = log_x_theta_student.to(torch.float64)
log_x_theta_teacher = log_x_theta_teacher.to(torch.float64)
if self.loss_type == "kl-fwd":
return (
log_x_theta_teacher.exp() * (log_x_theta_teacher - log_x_theta_student)
).sum(-1)
elif self.loss_type == "kl-bwd":
return (
log_x_theta_student.exp() * (log_x_theta_student - log_x_theta_teacher)
).sum(-1)
def training_step(self, batch, batch_idx):
self.log(
name="dt",
value=self._compute_dt(),
on_step=True,
on_epoch=False,
sync_dist=True,
)
return super().training_step(batch, batch_idx)
class CANDI(trainer_base.Diffusion):
def __init__(self, config, tokenizer):
self.mask_index = len(tokenizer)
vocab_size = len(tokenizer) + 1
super().__init__(config, tokenizer, vocab_size=vocab_size)
self._validate_configuration()
self.save_hyperparameters()
self.pure_continuous = config.algo.pure_continuous
self.continuous_noise =torch.linspace(.2, 4.0, 1000)
self.discrete_noise = expected_rank(d=self.vocab_size-1, a=torch.tensor(1).to(self.device), sigma=self.continuous_noise)
self.continuous_noise = torch.tensor(self.continuous_noise)
self.discrete_noise = torch.tensor(self.discrete_noise)
self.step_size = config.algo.step_size
self.min_percentile = config.algo.min_percentile
self.max_percentile = config.algo.max_percentile
self.sigma_min = config.algo.sigma_min
self.sigma_max = config.algo.sigma_max
self.temp = config.algo.temp
self.sampler = config.algo.sampler
self.is_embed = config.algo.is_embed
self.use_percentile_scheduling = config.algo.use_percentile_scheduling
def ve_sde_noise_sched(self, t, sigma_min=0.01, sigma_max=10.0):
return sigma_min * (sigma_max / sigma_min) ** t
def sample_t(self, n, accum_step):
return super()._sample_t(n, accum_step)
def get_continuous_from_discrete_noise(self, discrete_noise):
if self.is_embed:
sigmas = training_sigma_ve(discrete_noise, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
return sigmas
else:
target_percentile = discrete_noise * (self.max_percentile - self.min_percentile) + self.min_percentile
return sigma_from_time_vectorized(target_percentile, self.continuous_noise, self.discrete_noise)
# inference schedule
# while the embedding diffusion should use the VE schedule, this produced better results
# at least on text8 -- both schedules for embedding diffusion collapsed on OWT
def get_continuous_noise_sched_pure_cont(self, timesteps):
target_percentile = timesteps * (self.max_percentile - self.min_percentile) + self.min_percentile
return sigma_from_time_vectorized(target_percentile, self.continuous_noise, self.discrete_noise)
def to(self, *args, **kwargs):
super().to(*args, **kwargs)
self.continuous_noise = self.continuous_noise.to(*args, **kwargs)
self.discrete_noise = self.discrete_noise.to(*args, **kwargs)
return self
def _validate_configuration(self):
return
def _process_model_output(self, xt, model_output, reveal_mask):
if xt.ndim == 2:
xt_tokens = xt
else:
xt_tokens = xt.argmax(dim=-1)
model_output = model_output / self.temp
model_output = model_output - torch.logsumexp(
model_output, dim=-1, keepdim=True
)
reveal_mask = reveal_mask.bool()
# for pure continuous diffusion, reveal mask is all false -- so this doesnt have any impact
model_output[reveal_mask] = self.neg_infinity
model_output[reveal_mask, xt_tokens[reveal_mask]] = 0
return model_output
def nll_per_token(self,
log_x_theta,
x0_tokens,
alpha_t,
dalpha_t,
reveal_mask,
**kwargs):
log_p_theta = torch.gather(
input=log_x_theta, dim=-1, index=x0_tokens[:, :, None]
).squeeze(-1)
if self.pure_continuous:
nll = -1 * log_p_theta
else:
nll = log_p_theta * dalpha_t / (1 - alpha_t)
return nll
def discrete_noising(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
# gets noisy sample xt
def q_xt(self, x, alpha_t):
disc_xt = self.discrete_noising(x, alpha_t)
reveal_mask = (disc_xt == x).float()
discrete_noise = (1 - alpha_t).squeeze()
continuous_noise = self.get_continuous_from_discrete_noise(discrete_noise)
# accounting for the mask token
onehot = F.one_hot(x, num_classes=self.vocab_size-1).float()
if self.pure_continuous:
if self.is_embed:
clean_embeddings = self.backbone.get_embedding(x)
t = discrete_noise
sigma = self.get_continuous_from_discrete_noise(t).squeeze()
noise = torch.randn_like(clean_embeddings, device=x.device) * sigma[:, None, None]
xt = clean_embeddings + noise
# in pure continuous with embedding diffusion, we dont use any reveal mask
reveal_mask = torch.zeros_like(reveal_mask)
return {
'xt': xt,
'reveal_mask': reveal_mask,
'discrete_noise': discrete_noise,
'continuous_noise': continuous_noise,
'is_embed': True
}
else:
if self.cont_noise_schedule == 'corruption':
if self.sigma_table.device != discrete_noise.device:
self.sigma_table = self.sigma_table.to(discrete_noise.device)
self.error_table = self.error_table.to(discrete_noise.device)
sigma = error_to_sigma(discrete_noise, self.sigma_table, self.error_table)
else:
sigma = continuous_noise
noise = torch.randn_like(onehot, device=x.device) * sigma[:, None, None]
xt = onehot + noise
reveal_mask = torch.zeros_like(reveal_mask)
return {
'xt': xt,
'reveal_mask': reveal_mask,
'discrete_noise': discrete_noise,
'continuous_noise': continuous_noise,
'is_embed': False
}
xt_cont = onehot + continuous_noise[:, None, None].to(x.device) * torch.randn_like(onehot, device=x.device)
xt = onehot * reveal_mask.unsqueeze(-1) + (1 - reveal_mask).unsqueeze(-1) * xt_cont
return {
'xt': xt,
'reveal_mask': reveal_mask,
'discrete_noise': discrete_noise,
'continuous_noise': continuous_noise
}
def prior_sample(self, *batch_dims):
sigma = self.get_continuous_from_discrete_noise(torch.tensor(.999).to(self.device))
if self.is_embed:
sigma_max = self.sigma_max
noise = torch.randn(
*batch_dims, self.backbone.config.hidden_size, dtype=torch.float32, device=self.device) * sigma_max
else:
noise = torch.randn(
*batch_dims, self.vocab_size-1, dtype=torch.float32, device=self.device
) * sigma
return noise
def forward(self, **kwargs):
with torch.cuda.amp.autocast(dtype=torch.float32):
model_output = self.backbone(**kwargs)
return self._process_model_output(model_output=model_output, xt=kwargs['xt'], reveal_mask=kwargs['reveal_mask'])
def nll(self, x0, 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]
dalpha_t, alpha_t = self.noise(t)
alpha_t = alpha_t.unsqueeze(-1)
assert alpha_t.ndim == 2
noisy_output = self.q_xt(x0, alpha_t)
log_x_theta = self.forward(**noisy_output)
utils.print_nans(log_x_theta, "model_output")
return self.nll_per_token(
log_x_theta=log_x_theta,
alpha_t=alpha_t,
dalpha_t=dalpha_t,
x0_tokens=x0,
low_var=train_mode,
**noisy_output,
)
def generate_sample_prompt(self, prompt_tokens, x, prompt_mask):
prompt_onehots = torch.nn.functional.one_hot(prompt_tokens, x.size(-1)).float().to(x.device)
x = prompt_onehots * prompt_mask.unsqueeze(-1) + (1 - prompt_mask).unsqueeze(-1) * x
return x, prompt_mask
def generate_samples(self, **kwargs):
if self.pure_continuous or self.is_embed:
return self.generate_samples_pure_cont(**kwargs)
else:
if self.sampler == 'cached':
return self.generate_samples_cache(**kwargs)
else:
return self.generate_samples_nocache(**kwargs)
@torch.no_grad()
def generate_samples_nocache(self, num_samples, num_steps=None, eps=1e-5, prompt_tokens=None, prompt_mask=None):
if num_steps is None:
num_steps = self.config.sampling.steps
x = self.prior_sample(num_samples, self.num_tokens)
clean_mask = torch.zeros((num_samples, self.num_tokens), device=x.device)
if prompt_tokens is not None and prompt_mask is not None:
x, clean_mask = self.generate_sample_prompt(prompt_tokens, x, prompt_mask)
timesteps = torch.linspace(.999, eps, num_steps + 1, device=self.device)
if self.use_percentile_scheduling:
continuous_noise = self.get_continuous_from_discrete_noise(timesteps)
else:
continuous_noise = inference_sigmas(num_steps+1, self.sigma_min, self.sigma_max)
dt = (1 - eps) / (num_steps)
self.max_sigma = continuous_noise.max().item()
self.prev_px0 = None
for i in range(num_steps):
t = timesteps[i]
s = timesteps[i+1]
sigma_s = continuous_noise[i]
sigma_t = continuous_noise[i+1]
x_cont, p_x0 = self._continuous_step(x, t, sigma_s=sigma_s, sigma_t=sigma_t, clean_mask=clean_mask, time_s=s)
x, clean_mask = self._discrete_step(x_cont,
p_x0, t, dt, prev_clean_mask=clean_mask)
final_tokens = x.argmax(dim=-1)
return final_tokens
def generate_samples_pure_cont(self, num_samples, num_steps=None, eps=1e-5, prompt_tokens=None, prompt_mask=None):
if num_steps is None:
num_steps = self.config.sampling.steps
x = self.prior_sample(num_samples, self.num_tokens)
clean_mask = torch.zeros((num_samples, self.num_tokens), device=x.device)
if prompt_tokens is not None and prompt_mask is not None:
x, clean_mask = self.generate_sample_prompt(prompt_tokens, x, prompt_mask)
timesteps = torch.linspace(.999, eps, num_steps + 1, device=self.device)
continuous_noise = self.get_continuous_noise_sched_pure_cont(timesteps)
self.max_sigma = continuous_noise.max().item()
for i in range(num_steps):
t = timesteps[i]
s = timesteps[i+1]
sigma_s = continuous_noise[i]
sigma_t = continuous_noise[i+1]
new_x, denoised = self._continuous_step(x, t, sigma_s=sigma_s, sigma_t=sigma_t, clean_mask=clean_mask, time_s=s, is_embed=self.is_embed)
x = new_x
if self.is_embed:
final_tokens = denoised.argmax(dim=-1)
else:
final_tokens = x.argmax(dim=-1)
return final_tokens
# optimized function with cached embeddings
@torch.no_grad()
def generate_samples_cache(self, num_samples, num_steps=None, eps=1e-5, prompt_tokens=None, prompt_mask=None):
x = self.prior_sample(num_samples, self.num_tokens)
if num_steps is None:
num_steps = self.config.sampling.steps
embedding_cache = self.backbone.get_embedding(x)
clean_mask = torch.zeros((num_samples, self.num_tokens), device=x.device, dtype=torch.bool)
if prompt_tokens is not None and prompt_mask is not None:
x, clean_mask = self.generate_sample_prompt(prompt_tokens, x, prompt_mask)
timesteps = torch.linspace(.999, eps, num_steps + 1, device=self.device)
if self.use_percentile_scheduling:
continuous_noise = self.get_continuous_from_discrete_noise(timesteps)
else:
continuous_noise = inference_sigmas(num_steps+1, self.sigma_min, self.sigma_max)
dt = (1 - eps) / (num_steps)
self.max_sigma = continuous_noise.max().item()
# convert to tokens
x = x.argmax(dim=-1)
for i in range(num_steps):
t = timesteps[i]
s = timesteps[i+1]
sigma_s = continuous_noise[i]
sigma_t = continuous_noise[i+1]
embedding_cache, x0_hat = self._continuous_step_cache(x, t,
sigma_s=sigma_s,
sigma_t=sigma_t,
clean_mask=clean_mask.float(),
time_s=s,
embedding_cache=embedding_cache,
is_first=(i==0))
x, clean_mask = self._discrete_step_optimized(x0_hat, x, t, dt,
prev_clean_mask=clean_mask,
noise_removal_step=False)
return x
def _discrete_step_optimized(self, x0_hat, xt, t, dt, prev_clean_mask, noise_removal_step=False):
if noise_removal_step:
s=0
else:
s = t - dt
# unmasking correspends to 1-alpha(s) / 1-alpha(t)
# this is just s / t under a log linear schedule
unmask = torch.rand(prev_clean_mask.shape, device=prev_clean_mask.device) < (t-s)/t
xt[~prev_clean_mask] = x0_hat[~prev_clean_mask]
new_clean_mask = prev_clean_mask | unmask
return xt, new_clean_mask
def _discrete_step(self, x_sigma, p_x0, t, dt, prev_clean_mask, noise_removal_step=False):
if noise_removal_step:
s = 0.0
else:
s = t - dt
t_vec = torch.ones(x_sigma.shape[0], device=x_sigma.device) * t.item()
s_vec = torch.ones(x_sigma.shape[0], device=x_sigma.device) * s.item()
mask_probs = torch.ones((x_sigma.shape[0], x_sigma.shape[1], 1), device=x_sigma.device) * (s)
unmasked_probs = p_x0 * (t_vec - s_vec)[:, None, None]
q_xs = torch.cat([unmasked_probs, mask_probs], dim=-1)
_x = sample_categorical(q_xs)
new_clean_mask = (prev_clean_mask.bool() | (_x != self.mask_index)).float()
# For tokens that got sampled to real values (not mask), use those
old_x_tokens=x_sigma.argmax(dim=-1)
# For tokens that got sampled to mask, keep old tokens but mark as not clean
sampled_real_tokens = torch.where(_x != self.mask_index, _x, old_x_tokens)
# Apply copy logic: keep old tokens where prev_clean_mask is True
updated_tokens = torch.where(prev_clean_mask.bool(), old_x_tokens, sampled_real_tokens)
updated_x = torch.nn.functional.one_hot(updated_tokens, num_classes=x_sigma.shape[-1]).float().to(x_sigma.device)
updated_x = updated_x * new_clean_mask.unsqueeze(-1) + (1 - new_clean_mask).unsqueeze(-1) * x_sigma
return updated_x, new_clean_mask
@torch.no_grad()
def _continuous_step(
self,
x: torch.Tensor,
time_t: torch.Tensor,
time_s: torch.Tensor,
sigma_s: torch.Tensor,
sigma_t: torch.Tensor,
clean_mask: torch.Tensor=None,
is_embed=False,
) -> torch.Tensor:
dt_cont_vec = torch.ones(x.shape[0], device=x.device) * (sigma_s - sigma_t).item()
time_t_vec = torch.ones(x.shape[0], device=x.device) * time_t.item()
sigma_t_vec = torch.ones(x.shape[0], device=x.device) * sigma_t.item()
if clean_mask is None:
clean_mask = torch.zeros(x.shape[:-1], device=x.device)
cond_denoised = self.forward(xt=x, discrete_noise=time_t_vec, reveal_mask=clean_mask, continuous_noise=sigma_t_vec, is_embed=is_embed).double()
denoised = cond_denoised.exp()
if self.is_embed:
x0_hat = self.backbone.get_embedding(denoised)
else:
x0_hat = denoised
d = (x - x0_hat) / (sigma_t_vec[:, None, None] ** 2)
x_cont = x - dt_cont_vec[:, None, None] * d
return x_cont, denoised
def _continuous_step_cache(
self,
x: torch.Tensor,
time_t: torch.Tensor,
time_s: torch.Tensor,
sigma_s: torch.Tensor,
sigma_t: torch.Tensor,
embedding_cache: torch.Tensor,
clean_mask: torch.Tensor=None,
is_first=False,
) -> torch.Tensor:
dt = sigma_s - sigma_t
time_t_vec = torch.ones(x.shape[0], device=x.device) * time_t.item()
sigma_t_vec = torch.ones(x.shape[0], device=x.device) * sigma_t.item()
if clean_mask is None:
clean_mask = torch.zeros(x.shape[:-1], device=x.device)
cond_denoised = self.forward(xt=x,
discrete_noise=time_t_vec,
reveal_mask=clean_mask,
continuous_noise=sigma_t_vec,
embedding=embedding_cache).double()
denoised = cond_denoised.exp()
x0_hat = sample_categorical(denoised)
embedding_hat = self.backbone.get_embedding(x0_hat)
d = (embedding_cache - embedding_hat) / (sigma_t ** 2)
new_embedding_cache = embedding_cache - dt* d * self.step_size
return new_embedding_cache, x0_hat
def sigma_from_time_vectorized(
t: torch.Tensor, sigmas: torch.Tensor, errors: torch.Tensor
) -> torch.Tensor:
"""
Vectorized version of sigma_from_time using linear interpolation.
Args:
t: tensor of target Bayes error rates (ε ∈ [0, 1])
sigmas: 1D tensor of σ values
errors: 1D tensor of corresponding ε(σ)
Returns:
Tensor of interpolated σ values
"""
t = t.clamp(min=0.0, max=1.0)
t = t.to(sigmas.device)
# Normalize ε ∈ [0, 1] to indices in [0, len(errors)-1]
indices = torch.searchsorted(errors, t, right=True).clamp(1, len(errors) - 1)
i0 = indices - 1
i1 = indices
e0 = errors[i0]
e1 = errors[i1]
s0 = sigmas[i0]
s1 = sigmas[i1]
interp_t = (t - e0) / (e1 - e0 + 1e-8)
return s0 + interp_t * (s1 - s0)
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 expected_rank(d: int, a: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
"""
Vectorized expected rank of the true coordinate in a one-hot vector
corrupted by Gaussian noise.
Args:
d: dimension (number of coordinates)
a: amplitude(s) of the true coordinate, scalar or tensor
sigma: noise standard deviation(s), scalar or tensor
Returns:
Expected rank(s), same broadcasted shape as a and sigma
"""
# ensure tensors
a = torch.as_tensor(a, dtype=torch.float32)
sigma = torch.as_tensor(sigma, dtype=torch.float32)
# standard normal CDF in torch
Phi = torch.distributions.Normal(0.0, 1.0).cdf
# compute p = P(other beats true coordinate)
p = Phi(-a / (sigma * torch.sqrt(torch.tensor(2.0, dtype=torch.float32))))
return (d - 1) * p / d
# below is the noise schedule used for embedding diffusion -- follows the VE SDE schedule
def training_sigma_ve(t: torch.Tensor, sigma_min: float, sigma_max: float) -> torch.Tensor:
"""
Map training time t ∈ [0,1] to sigma via log-linear schedule.
σ(t) = σ_min * (σ_max / σ_min) ** t
Args:
t: tensor of shape [...] with values in [0,1]
sigma_min: minimum noise level
sigma_max: maximum noise level
Returns:
σ(t) with same shape as t
"""
t = torch.as_tensor(t, dtype=torch.float32)
log_ratio = torch.log(torch.tensor(sigma_max / sigma_min, dtype=torch.float32).to(t.device))
return sigma_min * torch.exp(t * log_ratio)
def inference_sigmas(n_steps: int, sigma_min: float, sigma_max: float) -> torch.Tensor:
"""
Log-linear schedule of sigmas for inference, decreasing from sigma_max to sigma_min.
σ_k = σ_max * (σ_min / σ_max) ** (k / (n_steps-1))
Args:
n_steps: number of inference steps
sigma_min: minimum noise level
sigma_max: maximum noise level
Returns:
Tensor of shape [n_steps], from σ_max down to σ_min
"""
ks = torch.linspace(0, 1, n_steps, dtype=torch.float32)
return sigma_max * (sigma_min / sigma_max) ** ks
def gaussian_to_disc_corruption(K: int, sigma: float) -> float:
normal = Normal(0., 1.)
s = torch.linspace(-4*sigma, 1+6*sigma, 1001)
phi = normal.log_prob((s-1)/sigma).exp() / sigma
cdf_power = normal.cdf(s/sigma).pow(K-1)
return 1.0 - torch.trapz(cdf_power * phi, s).item()
def build_error_to_sigma_schedule(
vocab_size: int,
sigma_range: tuple = (0.01, 1.0),
num_points: int = 500,
device: str = 'cpu'
) -> tuple[torch.Tensor, torch.Tensor]:
sigmas = torch.linspace(*sigma_range, steps=num_points, device=device)
errors = torch.tensor(
[gaussian_to_disc_corruption(vocab_size, s.item()) for s in sigmas],
device=device
)
return sigmas, errors
def error_to_sigma(
error_rates: torch.Tensor,
sigmas: torch.Tensor,
errors: torch.Tensor
) -> torch.Tensor:
error_rates = error_rates.clamp(min=0.0, max=1.0)
# Find bracketing indices
indices = torch.searchsorted(errors, error_rates, right=True).clamp(1, len(errors)-1)
i0 = indices - 1
i1 = indices
# Get bracketing values
e0 = errors[i0]
e1 = errors[i1]
s0 = sigmas[i0]
s1 = sigmas[i1]
# Linear interpolation
interp_t = (error_rates - e0) / (e1 - e0 + 1e-8)
return s0 + interp_t * (s1 - s0)
def sigma_to_error(
sigma_vals: torch.Tensor,
sigmas: torch.Tensor,
errors: torch.Tensor
) -> torch.Tensor:
sigma_vals = sigma_vals.clamp(min=sigmas[0], max=sigmas[-1])
# Find bracketing indices
indices = torch.searchsorted(sigmas, sigma_vals, right=True).clamp(1, len(sigmas)-1)
i0 = indices - 1
i1 = indices
# Get bracketing values
s0 = sigmas[i0]
s1 = sigmas[i1]
e0 = errors[i0]
e1 = errors[i1]
# Linear interpolation
interp_t = (sigma_vals - s0) / (s1 - s0 + 1e-8)
return e0 + interp_t * (e1 - e0)
import models
import os
import time
import pdb
import collections
import copy
import pickle
import fsspec
import numpy as np
import torch
import torch.nn.functional as F
from torch.distributions import Normal
import trainer_base
import utils
import torch
from omegaconf import DictConfig
# torch.serialization.add_safe_globals([DictConfig])
class AR(trainer_base.TrainerBase):
def __init__(self, config, tokenizer):
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
super().__init__(config, tokenizer, vocab_size=vocab_size)
self.save_hyperparameters()
self._validate_configuration()
self.temp = config.algo.get("temp", 1.0)
def _validate_configuration(self):
super()._validate_configuration()
assert not self.config.algo.time_conditioning
assert self.config.prior.type == "none"
def _process_model_input(self, x0, valid_tokens):
input_tokens = x0[:, :-1]
output_tokens = x0[:, 1:]
valid_tokens = valid_tokens[:, 1:]
return input_tokens, output_tokens, valid_tokens
def nll(self, input_tokens, output_tokens, current_accumulation_step, *args):
del current_accumulation_step
output = self.backbone(input_tokens, None)
output[:, :, self.mask_index] = self.neg_infinity
output = output.log_softmax(-1)
return -output.gather(-1, output_tokens[:, :, None])[:, :, 0]
def generate_samples(self, num_samples, prompt_mask=None, prompt_tokens=None, **kwargs):
# precompute token buffer
num_pred_tokens = self.num_tokens - 1
x = torch.zeros(
(num_samples, num_pred_tokens + 1), dtype=torch.long, device=self.device
)
if prompt_mask is not None and prompt_tokens is not None:
x = (prompt_mask * prompt_tokens + (1 - prompt_mask) * x).long()
start = prompt_mask.sum(1).max().long().item() - 1
else:
start = 0
x[:, 0] = self.tokenizer.bos_token_id
# precompute noise
noise = (
torch.distributions.Gumbel(0, 1)
.sample((num_samples, num_pred_tokens, self.vocab_size))
.to(self.device)
)
if self.config.sampling.use_float64:
noise = noise.to(torch.float64)
for i in range(start, num_pred_tokens):
output = self.backbone(x[:, : i + 1], None)
output[:, :, self.mask_index] = self.neg_infinity
output = output / self.temp
output = output.log_softmax(-1)
y = (output[:, -1, :] + noise[:, i, :]).argmax(-1)
x[:, i + 1] = y
return x
def _process_sigma(self, sigma):
del sigma
return None
class MDLM(trainer_base.AbsorbingState):
def __init__(self, config, tokenizer):
super().__init__(config, tokenizer)
self._validate_configuration()
def _validate_configuration(self):
# ancestral sampling isn't desirable because it's slow
assert self.sampler == "ancestral_cache" or self.sampler == "posterior"
def _process_model_output(self, model_output, xt, sigma):
del sigma
model_output = model_output / self.temp
model_output[:, :, self.mask_index] += self.neg_infinity