|
|
|
|
|
|
|
|
|
|
|
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) |
|
|
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) |
|
|
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 |
