File size: 1,063 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | import torch
from diffulex.sampler.auto_sampler import AutoSampler
from diffulex.sampler.base import DllmSamplerNoShiftBase
@AutoSampler.register("llada")
class LLaDASampler(DllmSamplerNoShiftBase):
def _compute_accepted_ids(
self,
block,
confidence: torch.Tensor,
initial_confidence: torch.Tensor,
sampled_tokens: torch.Tensor,
*,
threshold: float = 0.95,
**kwargs,
) -> torch.Tensor:
accept_threshold = block.thresholds.accept_threshold
high_conf_indices = torch.where(initial_confidence > accept_threshold)[0]
is_initial_block = block.block_id == 0 and block.prev_block is None
if block.should_force_decode_topk or is_initial_block:
topk_idx = (
torch.topk(confidence, 1)[1]
if len(high_conf_indices) == 0
else torch.tensor([], device=confidence.device, dtype=torch.long)
)
return torch.unique(torch.cat([topk_idx, high_conf_indices]))
return high_conf_indices
|