Caffin/bert-dllm

1"""Core routines for masked discrete language diffusion."""2 3from __future__ import annotations4 5from dataclasses import dataclass6from typing import Iterator, Literal, Protocol7 8import torch9 10 11SamplingStrategy = Literal["random", "confidence"]12 13 14class MaskedLanguageModel(Protocol):15    """Minimal ``transformers`` masked-LM protocol used by the sampler."""16 17    def __call__(self, *, input_ids: torch.Tensor, attention_mask: torch.Tensor): ...18 19 20@dataclass(frozen=True)21class DiffusionConfig:22    """The fixed canvas and reverse process used by the article reproduction."""23 24    canvas_length: int = 25625    prefix_length: int = 1626    denoising_steps: int = 1027 28    def __post_init__(self) -> None:29        if self.canvas_length <= self.prefix_length:30            raise ValueError("canvas_length must be greater than prefix_length")31        if self.denoising_steps < 1:32            raise ValueError("denoising_steps must be positive")33 34    @property35    def generated_length(self) -> int:36        """Number of tokens denoised after the fixed conditioning prefix."""37        return self.canvas_length - self.prefix_length38 39    @property40    def mask_probabilities(self) -> tuple[float, ...]:41        """Training mask rates, from fully masked to lightly masked."""42        return tuple(43            step / self.denoising_steps44            for step in range(self.denoising_steps, 0, -1)45        )46 47    def masks_after_step(self, step: int) -> int:48        """Return the exact number of canvas tokens to re-mask after one pass."""49        if not 1 <= step <= self.denoising_steps:50            raise ValueError("step is outside the denoising schedule")51        return round(self.generated_length * (self.denoising_steps - step) / self.denoising_steps)52 53 54@dataclass(frozen=True)55class DenoisingSnapshot:56    """A displayable state emitted after each denoising pass."""57 58    step: int59    input_ids: torch.Tensor60    mask_positions: torch.Tensor61    accepted_tokens: int62    total_tokens: int63    model_seconds: float64 65 66def prepare_conditioned_canvas(67    tokenizer: object,68    prompt: str,69    config: DiffusionConfig,70    device: torch.device,71) -> tuple[torch.Tensor, torch.Tensor, int]:72    """Build an all-mask canvas while preserving the article's fixed prefix.73 74    Short prompts are left-padded, exactly like the reference implementation.75    Long prompts are clipped rather than silently changing the conditioning length.76    """77    encoded = tokenizer(prompt, add_special_tokens=True, return_tensors="pt")78    prompt_ids = encoded["input_ids"].squeeze(0).to(dtype=torch.long)79    used_prompt_tokens = min(int(prompt_ids.numel()), config.prefix_length)80 81    if prompt_ids.numel() >= config.prefix_length:82        prefix = prompt_ids[: config.prefix_length]83    else:84        pad_id = getattr(tokenizer, "pad_token_id", None)85        if pad_id is None:86            raise ValueError("The tokenizer must define a pad_token_id")87        padding = torch.full(88            (config.prefix_length - prompt_ids.numel(),),89            int(pad_id),90            dtype=torch.long,91        )92        prefix = torch.cat((padding, prompt_ids))93 94    mask_id = getattr(tokenizer, "mask_token_id", None)95    if mask_id is None:96        raise ValueError("The tokenizer must define a mask_token_id")97 98    input_ids = torch.full(99        (1, config.canvas_length), int(mask_id), dtype=torch.long, device=device100    )101    input_ids[0, : config.prefix_length] = prefix.to(device)102    attention_mask = torch.ones_like(input_ids, device=device)103    return input_ids, attention_mask, used_prompt_tokens104 105 106def _choose_remask_positions(107    confidence: torch.Tensor,108    config: DiffusionConfig,109    target_count: int,110    strategy: SamplingStrategy,111    generator: torch.Generator,112) -> torch.Tensor:113    """Pick non-prefix positions to hide before the following reverse pass."""114    if target_count == 0:115        return torch.empty(0, dtype=torch.long, device=confidence.device)116 117    positions = torch.arange(118        config.prefix_length, config.canvas_length, device=confidence.device119    )120    if strategy == "confidence":121        return positions[torch.topk(confidence[positions], target_count, largest=False).indices]122    if strategy == "random":123        permutation = torch.randperm(124            positions.numel(), device=confidence.device, generator=generator125        )126        return positions[permutation[:target_count]]127    raise ValueError(f"Unsupported sampling strategy: {strategy}")128 129 130def denoise_canvas(131    model: MaskedLanguageModel,132    tokenizer: object,133    input_ids: torch.Tensor,134    attention_mask: torch.Tensor,135    config: DiffusionConfig,136    *,137    temperature: float,138    strategy: SamplingStrategy,139    generator: torch.Generator,140) -> Iterator[DenoisingSnapshot]:141    """Denoise a complete canvas in parallel and emit every reverse-process state.142 143    ``random`` reproduces the article's iterative re-masking. ``confidence`` retains144    the highest-confidence predictions, which is a common improved dLLM decoder.145    """146    if temperature <= 0:147        raise ValueError("temperature must be positive")148 149    mask_id = int(getattr(tokenizer, "mask_token_id"))150    blocked_ids = list(dict.fromkeys(getattr(tokenizer, "all_special_ids", [])))151    current = input_ids.clone()152    mask_positions = current.eq(mask_id)153    mask_positions[:, : config.prefix_length] = False154    model_seconds = 0.0155 156    for step in range(1, config.denoising_steps + 1):157        started_at = torch.cuda.Event(enable_timing=True) if current.is_cuda else None158        finished_at = torch.cuda.Event(enable_timing=True) if current.is_cuda else None159        if started_at is not None:160            started_at.record()161        with torch.inference_mode():162            logits = model(input_ids=current, attention_mask=attention_mask).logits163        if finished_at is not None:164            finished_at.record()165            finished_at.synchronize()166            model_seconds += started_at.elapsed_time(finished_at) / 1000167 168        logits = logits / temperature169        if blocked_ids:170            logits[..., blocked_ids] = -torch.inf171        probabilities = torch.softmax(logits, dim=-1)172        confidence = probabilities.max(dim=-1).values[0]173 174        active_positions = mask_positions[0]175        active_probabilities = probabilities[0, active_positions]176        sampled_tokens = torch.multinomial(177            active_probabilities, 1, generator=generator178        ).squeeze(-1)179        current[0, active_positions] = sampled_tokens180 181        target_count = config.masks_after_step(step)182        next_masked = _choose_remask_positions(183            confidence,184            config,185            target_count,186            strategy,187            generator,188        )189        mask_positions = torch.zeros_like(current, dtype=torch.bool)190        mask_positions[0, next_masked] = True191        current[mask_positions] = mask_id192        yield DenoisingSnapshot(193            step=step,194            input_ids=current.clone(),195            mask_positions=mask_positions.clone(),196            accepted_tokens=config.generated_length - target_count,197            total_tokens=config.generated_length,198            model_seconds=model_seconds,199        )200