Update dd_generator

dd_generator.py

@@ -0,0 +1,465 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+import torch.distributions as dists
+from torch.nn.utils.rnn import pad_sequence
+
+import numpy as np
+
+import math
+
+
+import sacrebleu
+
+from rouge import Rouge
+
+@dataclass
+class DiscreteDiffusionGeneratorArguments:
+    max_iterations: int = field(
+        default=10
+    )
+    mbr: int = field(
+        default=1
+    )
+    length_beam: int = field(
+        default=1
+    )
+    oracle_length: bool = field(
+        default=False
+    )
+    strategy: str = field(
+        default="reparam-uncond-deterministic-cosine"
+    )
+    argmax_decoding: bool = field(
+        default=True
+    )
+    bpe: str = field(
+        default="sentencepiece"
+    )
+    bleu_tokenize: str = field(
+        default="13a"
+    )
+    return_history: bool = field(
+        default=False
+    )
+    temperature: float = field(
+        default=0.8
+    )
+
+
+
+def topk_masking(scores, cutoff_len, stochastic=False, temp=1.0):
+    """
+    scores: [b, n]
+    cutoff_len: [b, 1]
+    stochastic: bool, whether to add noise to select top_k or not
+    returns:
+        mask: [b, n], with 1 if the token is in top-k lowest scores, 0 otherwise
+    """
+    if stochastic:
+        gumbel_noise = -torch.log(-torch.log(torch.rand_like(scores) + 1e-8) + 1e-8)
+        _scores = scores + temp * gumbel_noise
+    else:
+        _scores = scores
+    sorted_index = _scores.sort(-1)[0]
+    cutoff = sorted_index.gather(dim=-1, index=cutoff_len) # + 1e-10
+    # cutoff_len = k -> select k + 1 tokens
+    masking = _scores < cutoff
+    try:
+        assert (~(cutoff_len == 0).all()) | (~masking).all()
+    except:
+        import ipdb;ipdb.set_trace()
+    return masking
+
+
+class MergeBLEU(object):
+    def __call__(self, evalpreds):
+        # if torch.distributed.get_rank() == 0:
+        #     import ipdb; ipdb.set_trace()
+        # else:
+        #     import time; time.sleep(120)
+        import inspect
+        sys_stats, ref_stats = evalpreds[0], evalpreds[1]
+        
+        sys_stats = sys_stats.reshape(-1, 5).astype('long').sum(0).tolist()
+        ref_stats = ref_stats.reshape(-1, 5).astype('long').sum(0).tolist()
+        try:
+            from sacrebleu.metrics import BLEU
+            comp_bleu = BLEU.compute_bleu
+        except ImportError:
+            comp_bleu = sacrebleu.compute_bleu
+        fn_sig = inspect.getfullargspec(comp_bleu)[0]
+        if "smooth_method" in fn_sig:
+            smooth = {"smooth_method": "exp"}
+        else:
+            smooth = {"smooth": "exp"}
+        return {
+            "bleu": comp_bleu(
+                correct=sys_stats[:4], 
+                total=ref_stats[:4],
+                sys_len=sys_stats[-1],
+                ref_len=ref_stats[-1],
+                **smooth
+            ).score
+        }
+
+class MergeRouge(object):
+    def __call__(self, evalpreds):
+        # if torch.distributed.get_rank() == 0:
+        #     import ipdb; ipdb.set_trace()
+        # else:
+        #     import time; time.sleep(120)
+        import inspect
+        # sys
+        avg_rouge, batch_size = evalpreds[0], evalpreds[1]
+        
+        rouge = (avg_rouge * batch_size).sum() / batch_size.sum()
+        
+        return {
+            "rouge": rouge
+        }
+        
+
+class DiscreteDiffusionGenerator:
+    def __init__(self, args, dictionary=None, tokenizer=None) -> None:
+        self.args = args 
+        self.dictionary = dictionary
+        self.tokenizer = tokenizer
+        self.write_prediction = None
+    
+        assert (dictionary is not None) or (tokenizer is not None)
+        assert (dictionary is None) ^ (tokenizer is None)
+        
+        self.retain_history = args.return_history
+        
+        if dictionary is not None:
+            self.pad_id = dictionary.pad()
+            self.bos_id = dictionary.bos()
+            self.eos_id = dictionary.eos()
+            self.mask_id = dictionary.mask_index
+        else:
+            self.pad_id = tokenizer.pad_token_id
+            self.bos_id = tokenizer.bos_token_id
+            self.eos_id = tokenizer.eos_token_id
+            self.mask_id = tokenizer.mask_token_id
+    
+        self.rouge = Rouge(["rouge-l"])
+    
+    def set_write_to(self, path):
+        self.write_prediction = path
+    
+    def _reparam_decoding(
+        self, 
+        output_tokens, 
+        output_scores, 
+        cur_tokens,
+        cur_scores,
+        decoding_strategy,
+        xt_neq_x0, 
+        non_special_sym_mask, 
+        t,
+        max_step,
+        noise
+    ):
+        """
+            This function is used to perform reparameterized decoding.
+        """
+        # output_tokens: [B, N]
+        # output_scores: [B, N]
+        # cur_tokens: [B, N]
+        # cur_scores: [B, N]
+        # xt_neq_x0: equivalent to not_b_t [B, N]
+        # non_special_sym_mask: [B, N]
+        # noise: either [B, N] or scalar (if using the mask noise)
+        
+        # decoding_strategy needs to take the form of "reparam-<conditioning>-<topk_mode>-<schedule>"
+        _, condition, topk_mode, schedule = decoding_strategy.split("-")
+
+        # first set the denoising rate according to the schedule
+        if schedule == "linear":
+            rate = 1 - t / max_step
+        elif schedule == "cosine":
+            rate = np.cos(t / max_step * np.pi * 0.5)
+        else:
+            raise NotImplementedError
+
+        # compute the cutoff length for denoising top-k positions
+        cutoff_len = (
+            non_special_sym_mask.sum(1, keepdim=True).type_as(output_scores) * rate
+            ).long()
+        # set the scores of special symbols to a large value so that they will never be selected
+        _scores_for_topk = cur_scores.masked_fill(~non_special_sym_mask, 1000.0)
+        
+        # the top-k selection can be done in two ways: stochastic by injecting Gumbel noise or deterministic
+        if topk_mode.startswith("stochastic"):
+            noise_scale = float(topk_mode.replace("stochastic", ""))
+            lowest_k_mask = topk_masking(_scores_for_topk, cutoff_len, stochastic=True, temp=noise_scale * rate)
+        elif topk_mode == "deterministic":
+            lowest_k_mask = topk_masking(_scores_for_topk, cutoff_len, stochastic=False)
+        else:
+            raise NotImplementedError
+        
+        # Various choices to generate v_t := [v1_t, v2_t].
+        # Note that 
+        #   v1_t governs the outcomes of tokens where b_t = 1,
+        #   v2_t governs the outcomes of tokens where b_t = 0.
+        
+        # #### the `uncond` mode ####
+        # In our reparameterized decoding, 
+        # both v1_t and v2_t can be fully determined by the current token scores .
+        
+        # #### the `cond` mode ####
+        # However, we can also impose some conditional constraints on v1_t so that
+        # the decoding can be performed in a more conservative manner.
+        # For example, we can set v1_t = 0 only when 
+        # (the newly output tokens are the same as previous denoised results, AND
+        # the current token score becomes lower, AND
+        # the current token score is not in the top-k share among all tokens).
+        if condition == "cond":
+            not_v1_t = (cur_tokens == output_tokens) & (cur_scores < output_scores) & lowest_k_mask
+        elif condition == "uncond":
+            not_v1_t = lowest_k_mask
+        else:
+            raise NotImplementedError
+        
+        # for b_t = 0, the token is set to noise if it is in the lowest k scores.
+        not_v2_t = lowest_k_mask
+
+        masked_to_noise = (~xt_neq_x0 & not_v1_t) | (xt_neq_x0 & not_v2_t)
+        if isinstance(noise, torch.Tensor):
+            output_tokens.masked_scatter_(masked_to_noise, noise[masked_to_noise])
+        elif isinstance(noise, (int, float)):
+            output_tokens.masked_fill_(masked_to_noise, noise)
+        else:
+            raise NotImplementedError("noise should be either a tensor or a scalar")
+        output_scores.masked_fill_(masked_to_noise, -math.inf)
+
+        masked_to_x0 = xt_neq_x0 & ~not_v2_t
+        output_tokens.masked_scatter_(masked_to_x0, cur_tokens[masked_to_x0])
+        output_scores.masked_scatter_(masked_to_x0, cur_scores[masked_to_x0])
+        # b_{t} = (b_{t+1} & u_t) | v_t
+        # For convenience, save the NOT of b_t for the next iteration
+        # NOT_b_{t} = (NOT_b_{t+1} | not_v1_t) & not_v2_t
+        new_xt_neq_x0 = (xt_neq_x0 | not_v1_t) & not_v2_t
+        return new_xt_neq_x0
+    
+    def denoise_step(self, model, decoder_out, partial_masks):
+        output_tokens = decoder_out.output_tokens
+        output_scores = decoder_out.output_scores
+        prev_step, cur_step = decoder_out.step, decoder_out.step + 1 
+        max_step = decoder_out.max_step
+        temperature = self.args.temperature
+        # temperature = (
+        #     -0.05 * (cur_step / (max_step - 1)) + 0.5
+        #     if self.temperature_annealing
+        #     else self.temperature
+        # )
+        
+        # t = torch.LongTensor(
+        #     [(max_step - prev_step) * (model.num_diffusion_timesteps // max_step)] * output_tokens.size(0)
+        # ).to(output_tokens)
+        logits = model(output_tokens, partial_masks)
+        
+        logits[..., self.mask_id] = -math.inf
+        scores = torch.log_softmax(logits, dim=-1)
+        
+        
+        if self.args.strategy == "cmlm":
+            # get the mask
+            # <bos>, <eos> are ignored in this case since
+            # they are not equal to unk.
+            output_masks = output_tokens.eq(self.mask_id)
+            unmask_prob = 1 / (max_step - prev_step)
+            # where to unmask
+            changes = torch.rand(output_tokens.shape, device=output_tokens.device) < unmask_prob
+            # don't unmask somewhere already unmasked
+            changes = torch.bitwise_and(changes, output_masks)
+
+            if self.args.argmax_decoding:
+                output_scores, new_tokens = scores.max(-1)
+            else:
+                new_tokens = dists.Categorical(logits=scores / temperature).sample()
+                output_scores = torch.gather(scores, -1, new_tokens.unsqueeze(-1)).squeeze(-1)
+            output_tokens[changes] = new_tokens[changes]
+        elif self.args.strategy == "ar":
+            output_masks = output_tokens.eq(self.mask_id)
+            unmask_indices = (output_tokens.ne(self.mask_id) & output_tokens.ne(self.eos_id) & output_tokens.ne(self.pad_id)).sum(dim=-1)
+            indices = torch.arange(output_tokens.size(-1)).expand(output_tokens.shape).to(output_masks.device)
+            if self.args.argmax_decoding:
+                output_scores, new_tokens = scores.max(-1)
+            else:
+                new_tokens = dists.Categorical(logits=scores / temperature).sample()
+                output_scores = torch.gather(scores, -1, new_tokens.unsqueeze(-1)).squeeze(-1)
+            output_tokens[unmask_indices[:, None]==indices] = new_tokens[unmask_indices[:, None]==indices]
+            # output_tokens[changes] = new_tokens[changes]
+        else:
+            if self.args.argmax_decoding:
+                cur_scores, cur_tokens = scores.max(-1)
+            else:
+                cur_tokens = dists.Categorical(logits=scores / temperature).sample()
+                cur_scores = torch.gather(scores, -1, cur_tokens.unsqueeze(-1)).squeeze(-1)
+            cur_scores = cur_scores.to(output_scores)
+            
+            output_masks = self._reparam_decoding(
+                output_tokens=output_tokens,
+                output_scores=output_scores,
+                cur_tokens=cur_tokens,
+                cur_scores=cur_scores,
+                decoding_strategy=self.args.strategy,
+                xt_neq_x0=decoder_out.output_masks,
+                non_special_sym_mask=decoder_out.non_fixed_sym_masks,
+                t=cur_step,
+                max_step=max_step,
+                noise=self.mask_id
+            )
+        if self.retain_history:
+            history = ([] if decoder_out.history is None else decoder_out.history) + [output_tokens.clone()]
+        else:
+            history = None
+        # history = (
+        #     decoder_out.history + [output_tokens.clone()]
+        #     if self.retain_history
+        #     else None 
+        # )
+        return decoder_out._replace(
+            step=cur_step,
+            output_tokens=output_tokens,
+            output_scores=output_scores,
+            output_masks=output_masks,
+            history=history,
+        )
+
+    
+    def decode(self, seqs_tensors, preserve_special=False):
+        seqs_tensors[seqs_tensors < 0] = self.pad_id
+        if self.dictionary is not None:
+            seqs = [
+                self.dictionary.string(seq, self.args.bpe).strip()
+                for seq in seqs_tensors
+            ] 
+            if not preserve_special:
+                seqs = [seq.replace(self.dictionary.pad_word, '') for seq in seqs]
+        else:
+            seqs = self.tokenizer.batch_decode(seqs_tensors, skip_special_tokens=(not preserve_special))
+        return [seq.lower() for seq in seqs]
+    
+    def compute_bleu(self, hyps, refs):
+        if isinstance(hyps, torch.Tensor):
+            hyps = self.decode(hyps)
+        if isinstance(refs, torch.Tensor):
+            refs = self.decode(refs)
+        return sacrebleu.corpus_bleu(hyps, [refs], tokenize=self.args.bleu_tokenize)
+   
+    def compute_rouge(self, hyps, refs):
+        if isinstance(hyps, torch.Tensor):
+            hyps = self.decode(hyps)
+        if isinstance(refs, torch.Tensor):
+            refs = self.decode(refs) 
+        return self.rouge.get_scores(hyps, [[ref] for ref in refs])['rouge-l']['f']
+    
+    def stepwise_generate(self, model, inputs):
+        src_tokens = inputs["net_input"]["src_tokens"]
+        partial_masks = inputs["net_input"]["partial_masks"]
+        # assert src_tokens.size(-1) < 514
+        # assert partial_masks.size(-1) < 514
+        # target = inputs["target"]
+        raw_model = model.module if hasattr(model, "module") else model
+        if "prefix_masks" in inputs["net_input"]:
+            prefix_masks = inputs["net_input"]["prefix_masks"]
+        else:
+            prefix_masks = partial_masks
+        # TODO: FIXME: to support general blockwise generation.
+        partial_masks, prev_decoder_out = raw_model.initialize_decode_samples(
+            src_tokens, partial_masks, prefix_masks, oracle_length=self.args.oracle_length, length_beam=self.args.length_beam, mbr=self.args.mbr
+        )
+        prev_decoder_out = prev_decoder_out._replace(
+            step=0, max_step=self.args.max_iterations
+        )
+        for step in range(self.args.max_iterations):
+            prev_decoder_out = self.denoise_step(model, prev_decoder_out, partial_masks)
+            yield prev_decoder_out
+            
+    @torch.no_grad()
+    def generate(self, model, inputs):
+        src_tokens = inputs["net_input"]["src_tokens"]
+        partial_masks = inputs["net_input"]["partial_masks"]
+        # assert src_tokens.size(-1) < 514
+        # assert partial_masks.size(-1) < 514
+        # target = inputs["target"]
+        # TODO: FIXME: to support general blockwise generation.
+        if "prefix_masks" in inputs["net_input"]:
+            prefix_masks = inputs["net_input"]["prefix_masks"]
+        else:
+            prefix_masks = partial_masks 
+        partial_masks, prev_decoder_out = model.initialize_decode_samples(
+            src_tokens, partial_masks, prefix_masks, oracle_length=self.args.oracle_length, length_beam=self.args.length_beam, mbr=self.args.mbr
+        )
+        prev_decoder_out = prev_decoder_out._replace(
+            step=0, max_step=self.args.max_iterations
+        )
+        
+        for step in range(self.args.max_iterations):
+            prev_decoder_out = self.denoise_step(model, prev_decoder_out, partial_masks)            
+            
+        def finalized_hypos(tokens, scores, partial_mask, history=None):
+            cutoff = (
+                tokens.ne(self.pad_id) & 
+                tokens.ne(self.bos_id) & 
+                tokens.ne(self.eos_id) & 
+                (~partial_mask)
+            )
+            tokens = tokens[cutoff]
+            if scores is None:
+                score = None
+            else:
+                scores = scores[cutoff]
+                score = scores.mean().item()
+            ret_dict = {
+                "tokens": tokens,
+                "positional_scores": scores,
+                "score": score,
+                "alignment": None
+            }
+            if history is not None:
+                ret_dict["history"] = [
+                    finalized_hypos(history_tokens, None, partial_mask, history=None)
+                    for history_tokens in history
+                ]
+            return ret_dict
+        
+        def mbr_select(hyps):
+            index = np.argmax(np.array(
+                [self.rouge.get_scores([hyps[i]], [[hyps[j]]])['rouge-l']['f']
+                 for j in range(len(hyps)) if i != j]
+            ).mean() for i in range(len(hyps)))
+            return hyps[index]
+        
+        def score_select(hyps):
+            index = np.argmax([hyp["score"] for hyp in hyps])
+            return hyps[index]
+        
+        output_tokens, output_scores = prev_decoder_out.output_tokens, prev_decoder_out.output_scores
+        if self.retain_history:
+            full_history = prev_decoder_out.history 
+            histories = [[full_history[j][i] for j in range(self.args.max_iterations)] for i in range(output_tokens.size(0))]
+            hyps = []
+            for tokens, scores, partial_mask, history in zip(output_tokens, output_scores, partial_masks, histories):
+                hyps.append(finalized_hypos(tokens, scores, partial_mask, history))
+            # hyps = [
+            #     finalized_hypos(tokens, scores, partial_mask, history)
+            #     for tokens, scores, partial_mask, history in zip(output_tokens, output_scores, partial_masks, histories)
+            # ]
+        else:
+            hyps = [
+                finalized_hypos(tokens, scores, partial_mask, None) 
+                for tokens, scores, partial_mask in zip(output_tokens, output_scores, partial_masks)
+            ]
+        repeatition = self.args.mbr * self.args.length_beam
+        if repeatition > 1:
+            hyps = [score_select(hyps[i:i+repeatition])for i in range(0, len(hyps), repeatition)]
+            # hyps = [mbr_select(hyps[i:i+repeatition])for i in range(0, len(hyps), repeatition)]
+            
+        finalized = pad_sequence([h["tokens"] for h in hyps ], batch_first=True, padding_value=self.pad_id)
+        history = [[item["tokens"] for item in h["history"]] for h in hyps] if self.retain_history else None
+        return finalized, history