import torch import torch.nn.functional as F import numpy as np from tqdm import tqdm import utils import functools from tqdm.auto import tqdm @torch.no_grad() def sample_categorical(categorical_probs, temperature=1.0, dp=False): if temperature == 0.0: # Skip noise when temperature is 0 (sampling Gumbel is costly) return categorical_probs.argmax(dim=-1) noise = torch.rand_like(categorical_probs, dtype=( torch.float64 if dp else torch.float32)) gumbel_noise = (-torch.log(noise)) ** temperature return (categorical_probs / gumbel_noise).argmax(dim=-1) def safe_probability(p): if isinstance(p, torch.Tensor): p = torch.where(p < 0. and p > -1e-14, torch.zeros_like(p), p) p = torch.where( p > 1. and p < 1.000000001, torch.ones_like(p), p) return p else: if p < 0. and p > -1e-14: p = 0. if p > 1. and p < 1.000000001: p = 1. return p def torch_choice(x, num_tokens): batch_size = len(x) select_indices = [] for i in range(batch_size): chosen = x[i][torch.randperm( x[i].shape[-1])[:num_tokens]] select_indices.append(chosen) return torch.stack(select_indices) def _get_sampler(config, model, tokenizer, **kwargs): dict_samplers = { 'ar': ARSampler, 'analytic': AnalyticSampler, 'ancestral': AncestralCacheSampler, 'llada': LLaDaSampler, 'dimple': DimpleSampler, 'remdm-cap': ReMDMCap, 'remdm-rescale': ReMDMRescale, 'remdm-loop': ReMDMLoop, 'remdm-loop-conf': ReMDMLoop, 'forward-backward': ForwardBackward, 'llada-remdm-cap': LLaDaReMDM, 'llada-remdm-rescale': LLaDaReMDM, 'llada-remdm-loop': LLaDaReMDM, 'llada-remdm-loop-conf': LLaDaReMDM, # 'first_hitting': FirstHittingSampler, 'duo': DUOSampler, 'duo-llada': DUOLLaDaSampler, } if config.sampling.predictor in dict_samplers.keys(): return dict_samplers[config.sampling.predictor](config, model, tokenizer, **kwargs) else: raise ValueError( f"Invalid predictor: {config.sampling.predictor}") class Sampler(torch.nn.Module): def __init__(self, config, model, tokenizer): super().__init__() self.config = config self.model = model self.tokenizer = tokenizer self.p_nucleus = config.sampling.p_nucleus self.noise_removal = config.sampling.noise_removal self.use_float64 = config.sampling.use_float64 self.unmasking_temperature = config.sampling.temperature self.mask_index = model.mask_index self.num_tokens = model.num_tokens self.vocab_size = model.vocab_size self.device = model.device self.neg_infinity = torch.tensor( float('-inf'), device=self.device) self.backbone = model.backbone self.num_steps = self.config.sampling.steps self.gen_length = self.config.sampling.length if "length" in self.config.sampling else self.num_tokens def _initialize_x(self, num_samples, condition, target, target_idx): if condition is None: x = self.model.prior_sample( num_samples, self.gen_length) self.condition_mask = torch.zeros_like( x).type(torch.bool) else: # Match training length for entire sequence # x = self.model.prior_sample( # num_samples, gen_length) # Match training length for generated part max_condition_length = max( condition[n].shape[-1] for n in range(num_samples)) self.condition_mask = torch.ones( num_samples, self.gen_length + max_condition_length, device=self.device).type(torch.bool) x = [] for n in range(num_samples): _x = self.model.prior_sample( 1, self.gen_length) _x = torch.cat( [condition[n].to(_x.device), _x], dim=1) # _x = torch.nn.functional.pad( # _x, (0, max_condition_length - condition[n].shape[-1])) # Wait: this padding # (1) introduces artificial signal the Transformer will attend to # (2) I don't even remove it after generation # (3) completely messes up the -self.gen_length thing --> i need a padding mask assert num_samples == 1, "Batch sampling not supported for now because of padding issues" self.condition_mask[n, condition[n].shape[-1]: condition[n].shape[-1] + self.gen_length] = False x.append(_x) x = torch.cat(x, dim=0) if target is not None: x[:, : target_idx] = target return x @torch.no_grad() def forward(self, num_samples, eps, *args, logits_cache=None, target=None, target_idx=None, condition=None, **kwargs): """Generate samples from the model.""" x = self._initialize_x( num_samples, condition, target, target_idx) timesteps = torch.linspace( 1, eps, self.num_steps + 1, device=self.device) dt = (1 - eps) / self.num_steps conf = - \ torch.ones_like(x, device=self.device).to( torch.bfloat16) * torch.inf if hasattr(self.config.sampling, 'sample_latent') and self.config.sampling.sample_latent == "once": latent = self.model.sample_latent(x) else: latent = None # Iterative reverse process for i in tqdm(range(self.num_steps), desc='Sampling', disable=not getattr(self.config.sampling, 'verbose_progress', False)): t = timesteps[i] * \ torch.ones(x.shape[0], 1, device=self.device) if hasattr(self.config.sampling, 'sample_latent') and self.config.sampling.sample_latent == "each": latent = self.model.sample_latent(x) logits_cache, x, conf, latent = self._update( x=x, t=t, dt=dt, logits=logits_cache, conf=conf, condition=condition, latent=latent) if self.config.sampling.stop_when_eos: if (x[~self.condition_mask] == self.tokenizer.eos_token_id).sum() > 0: print("Found EOS, stopping") # return x[~(x == self.mask_index)] position_eos = (x[~self.condition_mask] == self.tokenizer.eos_token_id).nonzero() + self.condition_mask.sum() x = x[:position_eos[0]] break # Last step # if (x == self.mask_index).sum() > 0: t0 = timesteps[-1] * \ torch.ones(x.shape[0], 1, device=self.device) x = self._last_update(x=x, t0=t0, latent=latent) return x def _update(self, x, t, dt, latent=None, **kwargs): raise NotImplementedError def _last_update(self, x, t0, latent=None, **kwargs): if self.config.sampling.noise_removal == "greedy": _, alpha_t = self.model.noise(t0) sigma_t = self.model._sigma_from_alphat(alpha_t) x0 = self.model.forward_sample( x=x, sigma=sigma_t, latent=latent, prompt_index=self.condition_mask).argmax(dim=-1) elif self.config.sampling.noise_removal == "analytic": x0 = self._denoiser_update(x=x, t=t0, latent=latent) elif self.config.sampling.noise_removal == "analytic": x0 = self._denoiser_update(x=x, t=t0, latent=latent) elif self.config.sampling.noise_removal == "ancestral": _, alpha_t = self.model.noise(t0) sigma_t = self.model._sigma_from_alphat(alpha_t) alpha_s = torch.ones_like(alpha_t) logits = self.model.forward_sample( x, sigma_t, latent=latent, prompt_index=self.condition_mask) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum( sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) q_xs = p_x0 * (alpha_s - alpha_t)[:, :, None] q_xs[:, :, self.mask_index] = 1 - alpha_s _x = sample_categorical(q_xs, temperature=self.unmasking_temperature, dp=self.use_float64) copy_flag = (x != self.mask_index).to(x.dtype) x0 = copy_flag * x + (1 - copy_flag) * _x return x0 # Local model error proxies def indices_unmasking(self, x, x0, logits, num_tokens=None): if self.unmasking_strategy == 'random': all_tokens = torch.stack([torch.arange(x.shape[1]) for _ in range( x.shape[0])]).to(x.device) masked_tokens = [all_tokens[n, x[n] == self.mask_index] for n in range(x.shape[0])] return torch_choice(masked_tokens, num_tokens) probs = F.softmax(logits, dim=-1) if self.unmasking_strategy == 'confidence': x0_p = torch.squeeze( torch.gather(probs, dim=-1, index=torch.unsqueeze(x0, -1)), -1) x0 = torch.where(x == self.mask_index, x0, x) confidence = torch.where(x == self.mask_index, x0_p, -np.inf) criterion = confidence elif self.unmasking_strategy == 'margin': sorted_probs, _ = torch.sort(probs, dim=-1, descending=True) top1_probs = sorted_probs[:, 0] top2_probs = sorted_probs[:, 1] criterion = top1_probs - top2_probs elif self.unmasking_strategy == 'entropy': epsilon = 1e-10 log_probs = torch.log(probs + epsilon) criterion = torch.sum(probs * log_probs, dim=-1) else: raise NotImplementedError( f"Unmasking strategy {self.unmasking_strategy} not implemented") if num_tokens is None: # i.e. using dynamic number of tokens assert self.config.sampling.top_p_conf < 1. assert self.unmasking_strategy == "confidence", "Only confidence-based unmasking is supported with top-p for now" select_indices = [torch.arange(x.shape[1]).to(x.device)[ criterion[n] > self.config.sampling.top_p_conf] for n in range(x.shape[0])] else: if not hasattr(self, "token_temperature") or self.token_temperature == 0.: _, select_indices = torch.topk(criterion, k=num_tokens, dim=1) else: criterion = criterion / self.token_temperature criterion = F.softmax(criterion, dim=-1) select_indices = torch.multinomial(criterion, num_samples=num_tokens) return select_indices # Local model error proxies for remasking def compute_confidence_remdm(self, x, xs, p_x0, conf): unmask_mask = (x == self.mask_index) & ( xs != self.mask_index) batch_indices = torch.arange(xs.shape[0])[:, None].to(xs.device) feature_indices = torch.arange(xs.shape[1]).to(xs.device) conf_values = torch.squeeze( torch.gather(-p_x0, dim=-1, index=torch.unsqueeze(xs, -1)), -1) conf[unmask_mask] = conf_values[unmask_mask].to(conf.dtype) # Don't forget to put -inf if we have remasked remask_mask = (x != self.mask_index) & ( xs == self.mask_index) conf[remask_mask] = -np.inf return conf def compute_sigma_remdm(self, alpha_t, alpha_s, x, conf=None, logits=None, eps=1e-5): if self.remasking_strategy == "fixed": sigma = self.eta * torch.ones_like(x) elif self.remasking_strategy == "fixed-conf": sigma_max = self.eta eta = F.softmax( conf / self.remasking_temperature, dim=-1) if torch.isnan(eta).any(): # Should happen only at first step eta = torch.zeros_like(conf) masked_flag = (x == self.mask_index).to(torch.bool) eta[masked_flag] = 0 sigma = eta * sigma_max sigma = torch.min(sigma, torch.ones_like(x) * sigma_max / self.gen_length) # Assure that the bias we introduce with *=gen_length does not exceed the max sigma (and in pairtcular, not 1!!!) elif self.remasking_strategy == "fixed-margin": sigma_max = self.eta top12_tokens, _ = logits.topk(2, dim=-1) # should be positive (sorted in decreasing order) margin_tokens = top12_tokens[:, :, 0] - top12_tokens[:, :, 1] eta = F.softmax( margin_tokens / self.remasking_temperature, dim=-1) masked_flag = (x == self.mask_index).to(torch.bool) eta[masked_flag] = 0 sigma = eta * sigma_max elif self.remasking_strategy == "cap": sigma = torch.min(torch.ones_like(x) * self.eta, (1 - alpha_s) / alpha_t) if alpha_t > 0. else self.eta * torch.ones_like(x) elif self.remasking_strategy == "rescale": sigma_max = torch.min(torch.ones_like(x), (1 - alpha_s) / alpha_t) sigma = self.eta * sigma_max elif self.remasking_strategy == "rescale-conf": sigma_max = torch.min(torch.ones_like(x), (1 - alpha_s) / alpha_t) eta = F.softmax( conf / self.remasking_temperature, dim=-1) masked_flag = (x == self.mask_index).to(torch.bool) eta[masked_flag] = 0 sigma = eta * sigma_max else: raise NotImplementedError return sigma class ARSampler(Sampler): def _initialize_x(self, num_samples, condition, target, target_idx): noise = (torch.distributions.Gumbel(0, 1) .sample((num_samples, self.gen_length, self.vocab_size)) .to(self.device)) if condition is None: x = torch.zeros( (num_samples, self.gen_length), dtype=torch.long, device=self.device) else: max_condition_length = max( condition[n].shape[-1] for n in range(num_samples)) x = [] noise_complete = [] for n in range(num_samples): _x = torch.zeros( (1, self.gen_length), dtype=torch.long, device=self.device) _x = torch.cat( [condition[n].to(_x.device), _x], dim=1) _noise = torch.cat([ torch.zeros(condition[n].shape[-1], self.vocab_size).to(self.device), noise[n] ]).unsqueeze(0) assert num_samples == 1, "Batch sampling not supported for now because of padding issues" x.append(_x) noise_complete.append(_noise) x = torch.cat(x, dim=0) noise_complete = torch.cat(noise_complete, dim=0) if target is not None: x[:, : target_idx] = target return x, noise_complete @torch.no_grad() def forward(self, num_samples, eps, *args, logits_cache=None, target=None, target_idx=None, condition=None, **kwargs): """Generate samples from the model.""" x, noise = self._initialize_x( num_samples, condition, target, target_idx) if self.config.sampling.use_float64: noise = noise.to(torch.float64) kv_cache = self.config.sampling.kv_cache self.backbone.reset_kv_cache() sigma = torch.zeros(num_samples, dtype=self.model.dtype, device=self.device) i_start = min(condition[n].shape[-1] for n in range(num_samples)) -1 if condition is not None else 0 for i in range(i_start, i_start + self.gen_length - 2): output = self.backbone( x[:, :i + 1], sigma=sigma, x0=None, kv_cache=kv_cache) output[:, :, self.mask_index] = self.neg_infinity output = output.log_softmax(-1) if self.unmasking_temperature > 0: y = (output[:, -1, :] + noise[:, i, :] ** self.unmasking_temperature).argmax(-1) else: y = output[:, -1, :].argmax(-1) x[:, i + 1] = y # print("x tokens", x[0]) # print("x", self.tokenizer.decode(x[0])) # print("y tokens", y[0]) # print("y", y[0], self.tokenizer.decode(y[0])) # print("output tokens", output.argmax(dim=-1)[0]) # print("output", self.tokenizer.decode(output.argmax(dim=-1)[0])) self.backbone.reset_kv_cache() return x class AnalyticSampler(Sampler): 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 _update(self, x, t, dt, **kwargs): _, alpha_t = self.model.noise(t) _, alpha_s = self.model.noise(t - dt) sigma_t = self.model._sigma_from_alphat(alpha_t) sigma_s = self.model._sigma_from_alphat(alpha_s) dsigma = sigma_t - sigma_s score = self.model._get_score(x, sigma_t) if self.use_float64: score = score.to(torch.float64) stag_score = self._staggered_score(score, dsigma) probs = stag_score * \ self.model._transp_transition(x, dsigma) return sample_categorical(probs), None, None class AncestralCacheSampler(Sampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) assert self.unmasking_temperature > 0., "AncestralCacheSampler requires temperature > 0. because we don't sample from p_x0 but from q_xs, i.e. the mask_index will always be the largest logit" def _update(self, x, t, dt, logits=None, latent=None, **kwargs): _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma_t = self.model._sigma_from_alphat(alpha_t) if logits is None: if self.config.algo.name == "vdlm": if latent is None or self.config.sampling.sample_latent == "each": latent = self.model.sample_latent(x) else: latent = None logits = self.model.forward_sample( x, sigma_t, latent=latent) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus # always authorize at least the maximum-prob token top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) q_xs = p_x0 * (alpha_s - alpha_t)[:, :, None] q_xs[:, :, self.mask_index] = 1 - alpha_s _x = sample_categorical(q_xs, temperature=self.unmasking_temperature, dp=self.use_float64) copy_flag = (x != self.mask_index).to(x.dtype) xs = copy_flag * x + (1 - copy_flag) * _x logits_cache = logits if torch.allclose( xs, x) and not self.model.time_conditioning else None return logits_cache, xs, None, latent class LLaDaSampler(Sampler): """ Inspired from LLaDa generate method """ def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.num_tokens_per_step = self.gen_length // self.config.sampling.steps self.remaining_tokens = self.gen_length % self.config.sampling.steps self.unmasking_strategy = config.sampling.unmasking_strategy self.token_temperature = config.sampling.token_temperature def _update(self, x, t, dt, logits=None, latent=None, **kwargs): if hasattr(self.model, 'noise'): _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma_t = self.model._sigma_from_alphat(alpha_t) else: sigma_t = None logits = self.model.forward_sample( x, sigma=sigma_t, latent=latent, prompt_index=self.condition_mask) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus # always authorize at least the maximum-prob token top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) x0 = sample_categorical( p_x0, temperature=self.unmasking_temperature, dp=self.use_float64) # utils._save_tensor(x0, "x0_vdlm.pt" if latent is not None else "x0_mdlm.pt") num_tokens = self.remaining_tokens if t.sum() == 0 and ( x == self.mask_index).sum() > 0 else self.num_tokens_per_step select_indices = self.indices_unmasking( x, x0, logits, num_tokens) # print("select_indices", select_indices - self.condition_mask.sum()) transfer_index = torch.zeros_like( x0, dtype=torch.bool, device=x0.device) transfer_index.scatter_(1, select_indices, True) x[transfer_index] = x0[transfer_index] return logits, x, None, latent class DimpleSampler(Sampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.fallback_strategy = config.sampling.fallback_strategy # self.warmup_conf = config.sampling.warmup_conf self.adjust_dynamic_num_tokens = config.sampling.adjust_dynamic_num_tokens self.top_p_conf = config.sampling.top_p_conf self.max_new_tokens = config.sampling.max_new_tokens self.num_tokens_per_step = self.gen_length // self.config.sampling.steps self.unmasking_strategy = config.sampling.unmasking_strategy self.token_temperature = config.sampling.token_temperature @torch.no_grad() def forward(self, num_samples, eps, *args, logits_cache=None, target=None, target_idx=None, condition=None, **kwargs): """Generate samples from the model. Over-ride default forward because we need a dynamic number of steps""" if self.adjust_dynamic_num_tokens: raise ValueError return super().forward(num_samples, eps, *args, logits_cache=logits_cache, target=target, target_idx=target_idx, condition=condition, **kwargs) else: x = self._initialize_x( num_samples, condition, target, target_idx) while (x == self.mask_index).sum(): # Just for activating warm-up conf # t = (x == self.mask_index).sum() / self.gen_length _, x, _, _ = self._update( x=x, t=None, dt=None, logits=None, condition=condition) return x def _update(self, x, t, dt, logits=None, latent=None, **kwargs): if hasattr(self.model.backbone, 'adaLN') and self.model.backbone.adaLN: _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma_t = self.model._sigma_from_alphat(alpha_t) else: sigma_t = None logits = self.model.forward_sample( x, sigma=sigma_t, latent=latent, prompt_index=self.condition_mask) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus # always authorize at least the maximum-prob token top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) x0 = sample_categorical(p_x0, temperature=self.unmasking_temperature, dp=self.use_float64) select_indices = self.indices_unmasking( x, x0, logits, num_tokens=None) for j in range(x.shape[0]): _select_indices = select_indices[j][: self.max_new_tokens] if not len(_select_indices): if self.config.sampling.fallback_strategy == 'random': masked_ = (x[j] == self.mask_index) _select_indices = torch.randperm( masked_.sum(-1).item())[: self.num_tokens_per_step] _select_indices = masked_.nonzero()[ _select_indices].squeeze(1) elif self.config.sampling.fallback_strategy == 'greedy': probs_j = F.softmax(logits[j], dim=-1) x0_p_j = torch.squeeze( torch.gather(probs_j, dim=-1, index=torch.unsqueeze(x0, -1)), -1) x0_j = torch.where(x[j] == self.mask_index, x0[j], x[j]) confidence_j = torch.where(x[j] == self.mask_index, x0_p_j, -np.inf) _, _select_indices = torch.topk( confidence_j, k=1, dim=1) x[j, _select_indices] = x0[j, _select_indices] return None, x, None, None class ReMDMCap(Sampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.eta = config.sampling.eta self.copy_only_remasked = config.sampling.copy_only_remasked def _update(self, x, t, dt, logits=None, latent=None, condition=None, **kwargs): _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma_t = self.model._sigma_from_alphat(alpha_t) if logits is None: logits = self.model.forward_sample( x, sigma_t, latent=latent) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus # always authorize at least the maximum-prob token top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) xs = x[:, -self.gen_length:].clone() p_x0 = p_x0[:, -self.gen_length:] if alpha_t > 0: sigma = torch.min(torch.ones_like( alpha_t) * self.eta, (1 - alpha_s) / alpha_t) else: sigma = torch.ones_like(alpha_t) * self.eta q_xs = p_x0 * (1 - sigma)[:, :, None] q_xs[..., self.mask_index] = sigma[:, :, None] q_xs_2 = p_x0 * ((alpha_s - (1 - sigma) * alpha_t) / (1 - alpha_t))[:, :, None] q_xs_2[..., self.mask_index] = ( (1 - alpha_s - sigma * alpha_t) / (1 - alpha_t))[:, :, None] # First filter to discriminate between xt = M and xt ~= M copy_flag = (x[:, -self.gen_length:] != self.mask_index).to(torch.bool) q_xs = torch.where( copy_flag.unsqueeze(-1), q_xs, q_xs_2) xs = sample_categorical(q_xs, temperature=self.unmasking_temperature, dp=self.use_float64) if self.copy_only_remasked: # Second filter to keep the remasked tokens, but not sample from p_x0 again... (which is what sampling from q_xs does!!!) # This biases the ReMDM distribution but makes much more practical sense? copy_flag_2 = ( (x[:, -self.gen_length:] != self.mask_index) * (xs != self.mask_index) ).to(x.dtype) xs = copy_flag_2 * x[:, -self.gen_length: ] + (1 - copy_flag_2) * xs logits_cache = logits if torch.allclose( xs, x[:, -self.gen_length:]) and not self.model.time_conditioning else None x[:, -self.gen_length: ] = xs return logits_cache, x, None, None class ReMDMRescale(Sampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.eta = config.sampling.eta self.remasking_strategy = config.sampling.remasking_strategy self.remasking_temperature = config.sampling.remasking_temperature self.copy_only_remasked = config.sampling.copy_only_remasked def _update(self, x, t, dt, logits=None, latent=None, condition=None, conf=None, **kwargs): _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma_t = self.model._sigma_from_alphat(alpha_t) if logits is None: logits = self.model.forward_sample(x, sigma_t, latent=latent) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus # always authorize at least the maximum-prob token top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) xs = x[:, -self.gen_length:].clone() p_x0 = p_x0[:, -self.gen_length:] if (alpha_t > 0).all(): sigma_max = torch.min(torch.ones_like( alpha_t), (1 - alpha_s) / alpha_t) else: sigma_max = torch.ones_like(alpha_t) if "conf" in self.remasking_strategy: eta = F.softmax( conf[:, -self.gen_length:] / self.remasking_temperature, dim=-1) masked_flag = (x[:, -self.gen_length:] == self.mask_index).to(torch.bool) eta[masked_flag] = 0 elif self.remasking_strategy in ['none', 'random']: eta = self.eta * torch.ones_like(alpha_t) else: raise NotImplementedError( f"Remasking strategy {self.remasking_strategy} not implemented") sigma = eta * sigma_max q_xs = p_x0 * (1 - sigma)[:, :, None] q_xs[..., self.mask_index] = sigma q_xs_2 = p_x0 * ((alpha_s - (1 - sigma) * alpha_t) / (1 - alpha_t))[:, :, None] q_xs_2[..., self.mask_index] = ( (1 - alpha_s - sigma * alpha_t) / (1 - alpha_t)) copy_flag = (x[:, -self.gen_length:] != self.mask_index).to(torch.bool) q_xs = torch.where( copy_flag.unsqueeze(-1), q_xs, q_xs_2) xs = sample_categorical(q_xs, temperature=self.unmasking_temperature, dp=self.use_float64) if self.copy_only_remasked: # Second filter to keep the remasked tokens, but not sample from p_x0 again... (which is what sampling from q_xs does!!!) # This biases the ReMDM distribution but makes much more practical sense? copy_flag_2 = ( (x[:, -self.gen_length:] != self.mask_index) * (xs != self.mask_index) ).to(x.dtype) xs = copy_flag_2 * x[:, -self.gen_length: ] + (1 - copy_flag_2) * xs if "conf" in self.remasking_strategy: conf[:, -self.gen_length:] = self.compute_confidence_remdm( x[:, -self.gen_length:], xs, p_x0, conf[:, -self.gen_length:]) logits_cache = logits if torch.allclose( xs, x[:, -self.gen_length:]) and not self.model.time_conditioning else None x[:, -self.gen_length: ] = xs return logits_cache, x, conf, None class ReMDMLoop(Sampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.t_on = config.sampling.t_on self.t_off = config.sampling.t_off self.alpha_on = config.sampling.alpha_on self.eta = config.sampling.eta self.remasking_strategy = config.sampling.remasking_strategy self.remasking_temperature = config.sampling.remasking_temperature self.copy_only_remasked = config.sampling.copy_only_remasked def _update(self, x, t, dt, logits=None, latent=None, condition=None, conf=None, **kwargs): _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma_t = self.model._sigma_from_alphat(alpha_t) if logits is None: logits = self.model.forward_sample(x, sigma_t, latent=latent) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus # always authorize at least the maximum-prob token top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) xs = x[:, -self.gen_length:].clone() p_x0 = p_x0[:, -self.gen_length:] time = t[0].item() # use MDLM if time > self.t_on or time <= self.t_off: # compute alpha_t and alpha_s if time > self.t_on: move_chance_t = ( 1 - alpha_t * self.alpha_on / (1 - self.t_on)) move_chance_s = ( 1 - alpha_s * self.alpha_on / (1 - self.t_on)) elif time <= self.t_off: move_chance_t = ( (1 - alpha_t) * (1 - self.alpha_on) / self.t_off) move_chance_s = ( (1 - alpha_s) * (1 - self.alpha_on) / self.t_off) q_xs = p_x0 * (move_chance_t - move_chance_s)[:, :, None] q_xs[:, :, self.mask_index] = move_chance_s _x = sample_categorical(q_xs, temperature=self.unmasking_temperature, dp=self.use_float64) copy_flag = (x[:, -self.gen_length:] != self.mask_index).to(x.dtype) xs = copy_flag * x[:, -self.gen_length:] + (1 - copy_flag) * _x else: # use ReMDM sigma = self.compute_sigma_remdm( alpha_t, alpha_s, x[:, -self.gen_length:], conf=conf[:, -self.gen_length:], logits=logits[:, -self.gen_length:]) q_xs = p_x0 * (1 - sigma)[:, :, None] q_xs[..., self.mask_index] = sigma q_xs_2 = p_x0 * ((self.alpha_on - (1 - sigma) * self.alpha_on) / (1 - self.alpha_on))[:, :, None] q_xs_2[..., self.mask_index] = ( (1 - self.alpha_on - self.alpha_on * sigma) / (1 - self.alpha_on)) copy_flag = (x[:, -self.gen_length:] != self.mask_index).to(torch.bool) q_xs = torch.where( copy_flag.unsqueeze(-1), q_xs, q_xs_2) xs = sample_categorical(q_xs, temperature=self.unmasking_temperature, dp=self.use_float64) if self.copy_only_remasked: # Second filter to keep the remasked tokens, but not sample from p_x0 again... (which is what sampling from q_xs does!!!) # This biases the ReMDM distribution but makes much more practical sense? copy_flag_2 = ( (x[:, -self.gen_length:] != self.mask_index) * (xs != self.mask_index) ).to(x.dtype) xs = copy_flag_2 * x[:, -self.gen_length: ] + (1 - copy_flag_2) * xs if "conf" in self.remasking_strategy: conf[:, -self.gen_length:] = self.compute_confidence_remdm( x[:, -self.gen_length:], xs, p_x0, conf[:, -self.gen_length:]) logits_cache = logits if torch.allclose( xs, x[:, -self.gen_length:]) and not self.model.time_conditioning else None x[:, -self.gen_length: ] = xs return logits_cache, x, conf, None class LLaDaReMDM(Sampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.unmasking_strategy = config.sampling.unmasking_strategy self.eta = config.sampling.eta self.t_on = config.sampling.t_on self.t_off = config.sampling.t_off self.alpha_on = config.sampling.alpha_on self.remasking_strategy = config.sampling.remasking_strategy self.remasking_temperature = config.sampling.remasking_temperature self.remasked_tokens = 0 if "conf" in self.remasking_strategy: self.eta *= self.gen_length #To balance the fact that we are using a Bernoulli distribution, with weights approximately distributed as 1/gen_length through the softmax def _update(self, x, t, dt, logits=None, latent=None, condition=None, conf=None, **kwargs): self.num_tokens_per_step = (self.gen_length + self.remasked_tokens) // self.config.sampling.steps xs = x.clone() if hasattr(self.model.backbone, 'adaLN') and self.model.backbone.adaLN: _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma_t = self.model._sigma_from_alphat(alpha_t) else: sigma_t = None logits = self.model.forward_sample( x, sigma=sigma_t, latent=latent) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus # always authorize at least the maximum-prob token top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) x0 = sample_categorical( p_x0, temperature=self.unmasking_temperature, dp=self.use_float64) num_tokens = (x == self.mask_index).sum(dim=-1) if (t.sum(dim=-1) == 0).any() else self.num_tokens_per_step select_indices = self.indices_unmasking( xs, x0, logits, num_tokens) transfer_index = torch.zeros_like( x0, dtype=torch.bool, device=x0.device) transfer_index.scatter_(1, select_indices, True) xs[transfer_index] = x0[transfer_index] ### Remasking ### if self.remasking_strategy != "none" and t[0].item() <= self.t_on and t[0].item() > self.t_off: _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma = self.compute_sigma_remdm( alpha_t, alpha_s, xs[:, -self.gen_length:], conf=conf[:, -self.gen_length:], logits=logits[:, -self.gen_length:]) remasking_flag = torch.bernoulli(sigma).to(torch.bool) self.remasked_tokens += remasking_flag.sum()# remasking_flag.sum(dim=-1) #Works only for batch_size = 1 for now xs[:, -self.gen_length:][remasking_flag] = self.mask_index if "conf" in self.remasking_strategy: conf[:, -self.gen_length:] = self.compute_confidence_remdm( x[:, -self.gen_length:], xs[:, -self.gen_length:], p_x0[:, -self.gen_length:], conf[:, -self.gen_length:]) return None, xs, conf, None class ForwardBackward(Sampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.copy_only_remasked = config.sampling.copy_only_remasked def _update(self, x, t, dt, logits=None, latent=None, condition=None, conf=None, **kwargs): _, alpha_s = self.model.noise(t - dt) _, alpha_t = self.model.noise(t) sigma_t = self.model._sigma_from_alphat(alpha_t) if logits is None: logits = self.model.forward_sample(x, sigma_t, latent=latent) p_x0 = logits.exp() if self.p_nucleus < 1: sorted_probs, sorted_indices = torch.sort( p_x0, descending=True, dim=-1) cumulative_probs = torch.cumsum( sorted_probs, dim=-1) top_p_mask = cumulative_probs <= self.p_nucleus # always authorize at least the maximum-prob token top_p_mask[..., 0] = True nucleus_probs = sorted_probs * top_p_mask nucleus_probs /= nucleus_probs.sum( dim=-1, keepdim=True) p_x0 = torch.zeros_like( p_x0).scatter_(-1, sorted_indices, nucleus_probs) xs = x[:, -self.gen_length:].clone() p_x0 = p_x0[:, -self.gen_length:] if (alpha_t > 0).all(): sigma = (alpha_s - alpha_t) / alpha_t else: sigma = 1 q_xs = p_x0 * (1 - sigma)[:, :, None] q_xs[..., self.mask_index] = sigma q_xs_2 = p_x0 * \ ((alpha_s - (1 - sigma) * alpha_t) / (1 - alpha_t))[:, :, None] q_xs_2[..., self.mask_index] = ( (1 - alpha_s - sigma * alpha_t) / (1 - alpha_t)) copy_flag = (x[:, -self.gen_length:] != self.mask_index).to(torch.bool) q_xs = torch.where( copy_flag.unsqueeze(-1), q_xs, q_xs_2) xs = sample_categorical(q_xs, temperature=self.unmasking_temperature, dp=self.use_float64) if self.copy_only_remasked: # Second filter to keep the remasked tokens, but not sample from p_x0 again... (which is what sampling from q_xs does!!!) # This biases the ReMDM distribution but makes much more practical sense? copy_flag_2 = ( (x[:, -self.gen_length:] != self.mask_index) * (xs != self.mask_index) ).to(x.dtype) xs = copy_flag_2 * x[:, -self.gen_length: ] + (1 - copy_flag_2) * xs if "conf" in self.remasking_strategy: conf[:, -self.gen_length:] = self.compute_confidence_remdm( x[:, -self.gen_length:], xs, p_x0, conf[:, -self.gen_length:]) logits_cache = logits if torch.allclose( xs, x[:, -self.gen_length:]) and not self.model.time_conditioning else None x[:, -self.gen_length: ] = xs return logits_cache, x, conf, None def _get_esolm_sampler(config, model, tokenizer): if config.sampling.predictor == 'esolm_ancestral': # It is actually a first-hitting sampler return EsoLMAncestralSampler(config, model, tokenizer) if config.sampling.predictor == 'esolm_block_autoregressive': # It is actually a first-hitting sampler return EsoLMBlockARSampler(config, model, tokenizer) if config.sampling.predictor == 'esolm_block': # It is actually a first-hitting sampler return EsoLMBlockSampler(config, model, tokenizer) if config.sampling.predictor == 'esolm-confidence': return EsoLMConfidenceSampler(config, model, tokenizer) else: raise NotImplementedError( f"Sampler {config.sampling.predictor} not implemented") class EsoLMSampler(Sampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.kv_cache = config.sampling.kv_cache self.sequential_attn_mode = config.algo.sequential_attn_mode self.diffusion_attn_mode = config.algo.diffusion_attn_mode self.alpha_0 = config.sampling.alpha_0 if hasattr(config.sampling, 'alpha_0') and config.sampling.alpha_0 is not None else config.algo.alpha_0 def _get_esolm_sampler_config(self, num_steps, num_samples): raise NotImplementedError def _initialize_x(self, num_samples, condition, target, target_idx, sort_idx): noise = torch.distributions.Gumbel(0, 1).sample( (num_samples, self.gen_length, self.vocab_size)).to(self.device) if condition is None: x = self.model.prior_sample( num_samples, self.gen_length) self.condition_mask = torch.zeros_like( x).type(torch.bool) sort_idx_complete = sort_idx noise_complete = noise else: # Match training length for generated part max_condition_length = max( condition[n].shape[-1] for n in range(num_samples)) self.condition_mask = torch.ones( num_samples, self.gen_length + max_condition_length, device=self.device).type(torch.bool) x, sort_idx_complete, noise_complete = [], [], [] for n in range(num_samples): _x = self.model.prior_sample( 1, self.gen_length) if sort_idx is not None: _x = torch.gather(_x, dim=1, index=sort_idx[n].unsqueeze(0)) #Shift sort_idx to take into account the prompt _sort_idx_complete = torch.cat([ torch.arange(condition[n].shape[-1]).to(self.device), condition[n].shape[-1] + sort_idx[n] ], dim=-1).unsqueeze(0) _noise = torch.cat([ torch.zeros(condition[n].shape[-1], self.vocab_size).to(self.device), noise[n] ]).unsqueeze(0) _x = torch.cat( [condition[n].to(_x.device), _x], dim=1) assert num_samples == 1, "Batch sampling not supported for now because of padding issues" self.condition_mask[n, condition[n].shape[-1]: condition[n].shape[-1] + self.gen_length] = False x.append(_x) sort_idx_complete.append(_sort_idx_complete) noise_complete.append(_noise) x = torch.cat(x, dim=0) sort_idx_complete = torch.cat(sort_idx_complete, dim=0) noise_complete = torch.cat(noise_complete, dim=0) if target is not None: x[:, : target_idx] = target return x, sort_idx_complete, noise_complete @torch.no_grad() def forward(self, num_samples, eps, *args, target=None, target_idx=None, condition=None,**kwargs): """Generate samples from the model.""" (unmask_k_tokens, sort_idx) = self._get_esolm_sampler_config( num_steps=self.num_steps, num_samples=num_samples) num_diffusion_tokens = sum(unmask_k_tokens) assert num_diffusion_tokens == self.gen_length x, sort_idx, noise = self._initialize_x(num_samples, condition, target, target_idx, sort_idx) self.backbone.reset_kv_cache() if condition is not None: attn_mode = self.diffusion_attn_mode cutoffs = condition[0].shape[-1] _ = self.backbone.forward_sample( zt=x, # shape[1] is model.length sort_idx=sort_idx, # shape[1] is model.length attn_mode=attn_mode, cutoffs=cutoffs, kv_cache=self.kv_cache, last_k_start=0, curr_k_start=0, curr_k_end=condition[0].shape[-1]) unmasked_tokens = 0 if condition is None else condition[0].shape[-1] x = self.esolm_forward( x=x, sort_idx=sort_idx, unmask_k_tokens=unmask_k_tokens, noise=noise, num_diffusion_tokens=num_diffusion_tokens, unmasked_tokens=unmasked_tokens, ) self.backbone.reset_kv_cache() self.model.backbone.mdlm_mask = None self.model.backbone.seq_mask = None sort_idx_reversed = utils.get_reverse_indices(sort_idx) x = torch.gather(x, dim=1, index=sort_idx_reversed) return x @torch.no_grad() def esolm_forward(self, x, sort_idx, unmask_k_tokens, noise, num_diffusion_tokens, unmasked_tokens=0, last_k_start=None, return_indices=False, **ignored_kwargs): for i, k in enumerate(tqdm(unmask_k_tokens, desc='Sampling', disable=not getattr(self.config.sampling, 'verbose_progress', False))): if k: # attn_mode and cutoffs are important when kv caching is off if unmasked_tokens >= num_diffusion_tokens: # stop doing diffusion attn_mode = self.sequential_attn_mode # prefix-lm masking is named differently for diffusion # and sequential phase if attn_mode == 'mixed': # prefix-lm masking for sequential attn_mode = 'mixed2' # prefix-lm masking for diffusion cutoffs = num_diffusion_tokens else: # keep doing diffusion attn_mode = self.diffusion_attn_mode cutoffs = unmasked_tokens if i == 0: last_k_start = last_k_start if last_k_start is not None else 0 else: last_k_start = unmasked_tokens - \ unmask_k_tokens[i - 1] log_p_x0 = self.backbone.forward_sample( zt=x, # shape[1] is model.length sort_idx=sort_idx, # shape[1] is model.length attn_mode=attn_mode, cutoffs=cutoffs, kv_cache=self.kv_cache, last_k_start=last_k_start, curr_k_start=unmasked_tokens, # also last_k_end curr_k_end=unmasked_tokens + k) if self.use_float64: log_p_x0 = log_p_x0.to(torch.float64) log_p_x0[:, :, self.mask_index] = self.model.neg_infinity if self.p_nucleus < 1: log_p_x0 = utils.top_k_top_p_filtering( log_p_x0, top_p=self.p_nucleus) indices = slice( unmasked_tokens, unmasked_tokens + k) # print(indices, log_p_x0.shape, noise.shape, x.shape) if self.kv_cache: if self.unmasking_temperature > 0.: y = (log_p_x0 + 1 / self.unmasking_temperature * noise[:, indices, :]).argmax(-1) else: y = (log_p_x0).argmax(-1) else: if self.unmasking_temperature > 0.: y = (log_p_x0[:, indices, :] + 1 / self.unmasking_temperature * noise[:, indices, :]).argmax(-1) else: y = (log_p_x0[:, indices, :]).argmax(-1) x[:, indices] = y unmasked_tokens += k if return_indices: return x, unmasked_tokens - k, unmasked_tokens return x class EsoLMAncestralSampler(EsoLMSampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.unmasking_strategy = config.sampling.unmasking_strategy self.token_temperature = config.sampling.token_temperature def _get_esolm_sampler_config(self, num_steps, num_samples): remaining_tokens = self.gen_length unmask_k_tokens = [] dt = 1 / num_steps # Assumes a log-linear schedule. for t in np.linspace(1, dt, num_steps): _, alpha_t = self.model.noise(t) _, alpha_s = self.model.noise(t - dt) p_unmask = (alpha_s - alpha_t) / (1 - alpha_t) n_unmask = np.random.binomial( remaining_tokens, p_unmask) if n_unmask != 0: unmask_k_tokens.append(n_unmask) remaining_tokens -= n_unmask if remaining_tokens != 0 and self.alpha_0 == 1: unmask_k_tokens.append(remaining_tokens) num_diffusion_tokens = sum(unmask_k_tokens) sort_idx = torch.rand( num_samples, self.gen_length).argsort( descending=False).to(self.device) # Diffusion Tokens: shuffle # Sequential Tokens: don't shuffle sort_idx[:, num_diffusion_tokens:] = ( sort_idx[:, num_diffusion_tokens:].sort().values) unmask_k_tokens = unmask_k_tokens + [1] * ( self.gen_length - num_diffusion_tokens) return unmask_k_tokens, sort_idx class EsoLMBlockARSampler(EsoLMSampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.unmasking_strategy = config.sampling.unmasking_strategy self.token_temperature = config.sampling.token_temperature @functools.lru_cache def _get_esolm_sampler_config(self, num_steps, num_samples): sub_context_length = num_steps sar_steps = self.config.sampling.sar_steps num_parallel_decode = self.gen_length // ( self.config.sampling.sar_steps * num_steps) num_samples = num_samples assert num_parallel_decode, "num_parallel_decode must be greater than 0. Used num_steps: {}, sar_steps: {}, gen_length: {}, which results in num_parallel_decode: {}".format(num_steps, sar_steps, self.gen_length, num_parallel_decode) unmask_k_tokens = [] anchors = [] anchors = torch.arange( sar_steps * num_parallel_decode) * sub_context_length anchors = anchors.chunk(sar_steps) sort_idx = [] for sar_block_idx in range(sar_steps): unmask_k_tokens.extend( [1] * num_parallel_decode + [num_parallel_decode] * (sub_context_length - 1)) sort_idx.extend( [anchors[sar_block_idx] + i for i in range(sub_context_length)]) sort_idx = torch.cat(sort_idx).repeat(num_samples, 1).to(self.device) return unmask_k_tokens, sort_idx class EsoLMBlockSampler(EsoLMSampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.unmasking_strategy = config.sampling.unmasking_strategy self.token_temperature = config.sampling.token_temperature def _get_esolm_sampler_config(self, num_steps, num_samples): sub_context_length = num_steps num_parallel_decode = self.gen_length // num_steps num_samples = num_samples device = self.device sort_idx = torch.rand( num_samples, sub_context_length).argsort( descending=False).to(device) decoder_order = torch.arange( sub_context_length * num_parallel_decode).reshape( num_parallel_decode, sub_context_length).to(device) sort_idx = torch.vstack( [decoder_order[:, s].T.reshape(-1) for s in sort_idx]) unmask_k_tokens = [num_parallel_decode] * sub_context_length return unmask_k_tokens, sort_idx # class EsoLMConfidenceSampler(EsoLMSampler): # def __init__(self, config, model, tokenizer): # super().__init__(config, model, tokenizer) # self.unmasking_strategy = config.sampling.unmasking_strategy # self.token_temperature = config.sampling.token_temperature # @torch.no_grad() # def forward(self, num_samples, eps, *args, # target=None, # target_idx=None, # condition=None,**kwargs): # """Generate samples from the model.""" # # Lightning auto-casting is not working in this method for some reason # unmask_k_tokens = self._tokens_unmasked_per_step(self.num_steps) # num_diffusion_tokens = sum(unmask_k_tokens) # sort_idx = torch.rand( # num_samples, self.gen_length).argsort( # descending=False).to(self.device) # sort_idx[:, num_diffusion_tokens:] = ( # sort_idx[:, num_diffusion_tokens:].sort().values) # ### Here do something to pre-shuffle according to confidence?? # x, sort_idx, noise = self._initialize_x(num_samples, condition, target, target_idx, sort_idx) # unmask_k_tokens = unmask_k_tokens + [1] * ( # self.gen_length - num_diffusion_tokens) # assert sum(unmask_k_tokens) == self.gen_length # self.backbone.reset_kv_cache() # if condition is not None: # attn_mode = self.diffusion_attn_mode # cutoffs = condition[0].shape[-1] # _ = self.backbone.forward_sample( # zt=x, # shape[1] is model.length # sort_idx=sort_idx, # shape[1] is model.length # attn_mode=attn_mode, # cutoffs=cutoffs, # kv_cache=self.kv_cache, # last_k_start=0, # curr_k_start=0, # curr_k_end=condition[0].shape[-1]) # unmasked_tokens = 0 if condition is None else condition[0].shape[-1] # x = self.esolm_forward( # x=x, # sort_idx=sort_idx, # unmask_k_tokens=unmask_k_tokens, # noise=noise, # num_diffusion_tokens=num_diffusion_tokens, # unmasked_tokens=unmasked_tokens, # ) # self.backbone.reset_kv_cache() # self.model.backbone.mdlm_mask = None # self.model.backbone.seq_mask = None # sort_idx_reversed = utils.get_reverse_indices(sort_idx) # x = torch.gather(x, dim=1, index=sort_idx_reversed) # return x # @torch.no_grad() # def esolm_forward(self, x, sort_idx, unmask_k_tokens, noise, num_diffusion_tokens, unmasked_tokens=0, last_k_start=None, return_indices=False, **ignored_kwargs): # for i, k in enumerate(unmask_k_tokens): # if k: # # attn_mode and cutoffs are important when kv caching is off # if unmasked_tokens >= num_diffusion_tokens: # # stop doing diffusion # attn_mode = self.sequential_attn_mode # # prefix-lm masking is named differently for diffusion # # and sequential phase # if attn_mode == 'mixed': # prefix-lm masking for sequential # attn_mode = 'mixed2' # prefix-lm masking for diffusion # cutoffs = num_diffusion_tokens # else: # # keep doing diffusion # attn_mode = self.diffusion_attn_mode # cutoffs = unmasked_tokens # if i == 0: # last_k_start = last_k_start if last_k_start is not None else 0 # else: # last_k_start = unmasked_tokens - \ # unmask_k_tokens[i - 1] # log_p_x0 = self.backbone.forward_sample( # zt=x, # shape[1] is model.length # sort_idx=sort_idx, # shape[1] is model.length # attn_mode=attn_mode, # cutoffs=cutoffs, # kv_cache=self.kv_cache, # last_k_start=last_k_start, # curr_k_start=unmasked_tokens, # also last_k_end # curr_k_end=unmasked_tokens + k) # if self.use_float64: # log_p_x0 = log_p_x0.to(torch.float64) # log_p_x0[:, :, # self.mask_index] = self.model.neg_infinity # if self.p_nucleus < 1: # log_p_x0 = utils.top_k_top_p_filtering( # log_p_x0, top_p=self.p_nucleus) # indices = slice( # unmasked_tokens, unmasked_tokens + k) # # print(indices, log_p_x0.shape, noise.shape, x.shape) # if self.kv_cache: # if self.unmasking_temperature > 0.: # y = (log_p_x0 + 1 / self.unmasking_temperature * noise[:, indices, :]).argmax(-1) # else: # y = (log_p_x0).argmax(-1) # else: # if self.unmasking_temperature > 0.: # y = (log_p_x0[:, indices, :] + # 1 / self.unmasking_temperature * noise[:, indices, :]).argmax(-1) # else: # y = (log_p_x0[:, indices, :]).argmax(-1) # x[:, indices] = y # unmasked_tokens += k # if return_indices: # return x, unmasked_tokens - k, unmasked_tokens # return x class DUOSampler(Sampler): 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.model.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))) @torch.no_grad() def _update(self, x, t, dt, logits=None, conf=None, condition=None, latent=None): # if noise_removal_step: # alpha_s = torch.ones_like(alpha_t) # else: # _, alpha_s = self.noise(t - dt) _, alpha_t = self.model.noise(t) _, alpha_s = self.model.noise(t - dt) sigma_t = self.model._sigma_from_alphat(alpha_t) assert alpha_t.ndim == 2 x_theta = self.model(x, sigma_t, prompt_index=self.condition_mask).exp() q_xs = self._compute_posterior( x=x_theta, xt=x, alpha_s=alpha_s, alpha_t=alpha_t) if self.p_nucleus < 1: log_p_x0 = utils.top_k_top_p_filtering( q_xs.log(), top_p=self.p_nucleus) noise = torch.distributions.Gumbel(0, 1).sample( log_p_x0.shape).to(self.device) return None, (log_p_x0 + noise).argmax(dim=-1), None, None xs = sample_categorical(q_xs, temperature=self.unmasking_temperature, dp=self.use_float64) if condition is not None: #TMP a bit ugly xs[:, : condition[0].shape[-1]] = condition[0] return None, xs, None, None class DUOLLaDaSampler(DUOSampler): def __init__(self, config, model, tokenizer): super().__init__(config, model, tokenizer) self.num_tokens_per_step = self.gen_length // self.config.sampling.steps self.remaining_tokens = self.gen_length % self.config.sampling.steps self.unmasking_strategy = config.sampling.unmasking_strategy self.token_temperature = config.sampling.token_temperature def indices_unmasking(self, x, x0, logits, num_tokens=None, condition=None): # Need to adapt because of non-masking process self.condition_mask = self.condition_mask.type(torch.bool) if self.unmasking_strategy == 'random': all_tokens = torch.arange(x.shape[1]).unsqueeze(0).repeat(x.shape[0], 1).to(x.device) legit_tokens = [all_tokens[n, ~self.condition_mask[n]] for n in range(x.shape[0])] return torch_choice(legit_tokens, num_tokens) probs = F.softmax(logits, dim=-1) if self.unmasking_strategy == 'confidence': x0_p = torch.squeeze( torch.gather(probs, dim=-1, index=torch.unsqueeze(x0, -1)), -1) x0 = torch.where(self.condition_mask, x, x0) confidence = torch.where(self.condition_mask, -np.inf, x0_p) criterion = confidence else: raise NotImplementedError( f"Unmasking strategy {self.unmasking_strategy} not implemented") if not hasattr(self, "token_temperature") or self.token_temperature == 0.: _, select_indices = torch.topk(criterion, k=num_tokens, dim=1) else: criterion = criterion / self.token_temperature criterion = F.softmax(criterion, dim=-1) select_indices = torch.multinomial(criterion, num_samples=num_tokens) return select_indices @torch.no_grad() def _update(self, x, t, dt, logits=None, conf=None, condition=None, latent=None): _, alpha_t = self.model.noise(t) _, alpha_0 = self.model.noise(torch.zeros_like(t)) sigma_t = self.model._sigma_from_alphat(alpha_t) assert alpha_t.ndim == 2 x_theta = self.model(x, sigma_t, prompt_index=self.condition_mask).exp() q_x0 = self._compute_posterior( x=x_theta, xt=x, alpha_s=alpha_0, alpha_t=alpha_t) if self.p_nucleus < 1: log_p_x0 = utils.top_k_top_p_filtering( q_x0.log(), top_p=self.p_nucleus) noise = torch.distributions.Gumbel(0, 1).sample( log_p_x0.shape).to(self.device) x0 = (log_p_x0 + noise).argmax(dim=-1) else: x0 = sample_categorical(q_x0, temperature=self.unmasking_temperature, dp=self.use_float64) num_tokens = self.remaining_tokens if t.sum().item() == 0. else self.num_tokens_per_step select_indices = self.indices_unmasking( x, x0, q_x0.log(), num_tokens=num_tokens, condition=condition) transfer_index = torch.zeros_like( x0, dtype=torch.bool, device=x0.device) transfer_index.scatter_(1, select_indices, True) x[transfer_index] = x0[transfer_index] return None, x, None, None