Upload generation_utils.py with huggingface_hub

generation_utils.py

@@ -0,0 +1,133 @@
+ # Copyright (c) 2026 ByteDance Ltd. and/or its affiliates
+ # SPDX-License-Identifier: MIT
+
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
+from transformers.cache_utils import DynamicCache
+
+def add_gumbel_noise(logits, temperature):
+    if temperature == 0:
+        return logits
+    logits = logits.to(torch.float64)
+    noise = torch.rand_like(logits, dtype=torch.float64)
+    gumbel_noise = (- torch.log(noise)) ** temperature
+    return logits.exp() / gumbel_noise
+
+
+def get_num_transfer_tokens(mask_index, steps):
+    mask_num = mask_index.sum(dim=1, keepdim=True)
+
+    base = mask_num // steps
+    remainder = mask_num % steps
+
+    num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
+
+    for i in range(mask_num.size(0)):
+        num_transfer_tokens[i, :remainder[i]] += 1
+
+    return num_transfer_tokens
+
+def make_block_causal_mask(seq_len, block_size=2, device=None, dtype=torch.bool):
+    num_blocks = (seq_len + block_size - 1) // block_size
+    block_mask = torch.tril(torch.ones((num_blocks, num_blocks), dtype=torch.bool, device=device))
+    local_block = torch.ones((block_size, block_size), dtype=torch.bool, device=device)
+    mask = torch.kron(block_mask, local_block)[:seq_len, :seq_len]
+
+    attention_mask = mask.float()
+    attention_mask.masked_fill_(~mask, float('-inf'))
+    attention_mask = attention_mask.unsqueeze(0).unsqueeze(0).to(dtype)
+    return attention_mask
+
+@ torch.no_grad()
+def generate_block(model, prompt, steps=128, gen_length=128, block_length=128, temperature=0.,
+            remasking='low_confidence', tokenizer=None, mask_id=5, threshold=0.95, shift=False, eos_id=None):
+    x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(model.device)
+    x[:, :prompt.shape[1]] = prompt.clone()
+
+    assert gen_length % block_length == 0
+    num_blocks = gen_length // block_length
+
+    assert steps % num_blocks == 0
+    steps = steps // num_blocks
+
+    prompt_len = prompt.shape[1]
+    res_block = block_length - prompt_len % block_length
+    every_block = [block_length for _ in range(num_blocks)]
+    if res_block > 0:
+        every_block = [res_block] + every_block
+        every_block[-1] = block_length - res_block
+    cum_block = [sum(every_block[:i+1]) for i in range(len(every_block))]
+    num_block = len(cum_block)
+
+    block_diffusion_attention_mask = make_block_causal_mask(prompt.shape[1] + gen_length, block_length, model.device, dtype=torch.bfloat16)
+    nfe = 0  
+    final_flag = 0
+    prefill_length = prompt_len // block_length * block_length
+    if prefill_length > 0:
+        cur_attn_mask = block_diffusion_attention_mask[:, :, :prefill_length, :prefill_length]
+        past_key_values = model(x[:, :prefill_length], attention_mask=cur_attn_mask, use_cache=True).past_key_values
+    for num_block in range(num_blocks):
+        current_block_start = prompt_len + cum_block[num_block - 1] if num_block > 0 else prefill_length
+        current_block_end = prompt_len + cum_block[num_block]
+
+        block_mask_index = (x[:, current_block_start:current_block_end] == mask_id)
+        num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps)
+
+        replace_position = torch.zeros_like(x, dtype=torch.bool)
+        replace_position[:, current_block_start:current_block_end] = 1
+        i = 0
+        while True:
+            nfe += 1
+            mask_index = (x[:, current_block_start:current_block_end] == mask_id)
+            cur_attn_mask = block_diffusion_attention_mask[:, :, current_block_start:current_block_end, :current_block_end]
+            output = model(x[:, current_block_start:current_block_end], attention_mask=block_diffusion_attention_mask[:, :, current_block_start:current_block_end, :current_block_end], past_key_values=past_key_values, use_cache=True, cache_position=replace_position.nonzero(as_tuple=True)[1])
+            logits = output.logits
+            x0, transfer_index = get_transfer_index(logits, temperature, remasking, mask_index, 
+                                            x[:, current_block_start:current_block_end], num_transfer_tokens[:, i] if threshold is None else None, threshold, shift=False)
+            x[:, current_block_start:current_block_end][transfer_index] = x0[transfer_index]
+            if (x[:, current_block_start:current_block_end] == mask_id).sum() == 0:
+                if eos_id is not None and (x[:, current_block_start:current_block_end] == eos_id).sum() > 0:
+                    final_flag = 1
+                    x = x[:, :current_block_end]
+                    break
+                past_key_values = model(x[:, current_block_start:current_block_end], attention_mask=block_diffusion_attention_mask[:, :, current_block_start:current_block_end, :current_block_end], past_key_values=past_key_values, use_cache=True, cache_position=replace_position.nonzero(as_tuple=True)[1]).past_key_values
+                break
+        if final_flag == 1:
+            break
+            i += 1
+    return x, nfe
+
+
+def get_transfer_index(logits, temperature, remasking, mask_index, x, num_transfer_tokens, threshold=None, shift=False):
+    logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
+    x0 = torch.argmax(logits_with_noise, dim=-1) # b, l
+    if shift == True:
+        x0 = torch.cat([x[:, :1], x0[:, :-1]], dim=-1)
+        pad = torch.zeros_like(logits[:, :1])
+        logits = torch.cat([pad, logits[:, :-1]], dim=1)
+    if remasking == 'low_confidence':
+        p = F.softmax(logits.to(torch.float64), dim=-1)
+        x0_p = torch.squeeze(
+            torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) # b, l
+    elif remasking == 'random':
+        x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
+    else:
+        raise NotImplementedError(remasking)
+    
+    x0 = torch.where(mask_index, x0, x)
+    confidence = torch.where(mask_index, x0_p, -np.inf)
+
+    transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
+    if threshold is not None:
+        num_transfer_tokens = mask_index.sum(dim=1, keepdim=True)
+    for j in range(confidence.shape[0]):
+        _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j])
+        transfer_index[j, select_index] = True
+        if threshold is not None:
+            for k in range(1, num_transfer_tokens[j]):
+                if confidence[j, select_index[k]] < threshold:
+                    transfer_index[j, select_index[k]] = False
+    return x0, transfer_index