| """Student scaffold for absorbing-mask diffusion sampling.""" |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import Tensor |
|
|
| |
| |
| from model import AnonymousEncoder |
|
|
|
|
| @dataclass |
| class DiffusionStepOutput: |
| input_ids: Tensor |
| sampled_ids: Tensor |
| confidence: Tensor |
| remaining_masks: Tensor |
|
|
|
|
| @dataclass |
| class DiffusionOutput: |
| input_ids: Tensor |
| history: tuple[Tensor, ...] |
|
|
|
|
| def cosine_mask_ratio(step: int, total_steps: int) -> float: |
| """Fraction of editable positions that remain masked after this step.""" |
| raise NotImplementedError |
|
|
|
|
| def corrupt( |
| clean_ids: Tensor, |
| editable_mask: Tensor, |
| *, |
| step: int, |
| total_steps: int, |
| mask_id: int, |
| generator: torch.Generator | None = None, |
| ) -> Tensor: |
| """Sample x_t from the absorbing-mask forward process q(x_t | x_0).""" |
| raise NotImplementedError |
|
|
|
|
| def diffusion_step( |
| model: AnonymousEncoder, |
| input_ids: Tensor, |
| editable_mask: Tensor, |
| *, |
| step: int, |
| total_steps: int, |
| mask_id: int, |
| pad_id: int, |
| temperature: float = 1.0, |
| top_k: int | None = None, |
| forbidden_token_ids: tuple[int, ...] = (), |
| generator: torch.Generator | None = None, |
| ) -> DiffusionStepOutput: |
| """Perform one confidence-based reverse diffusion transition x_t -> x_(t-1).""" |
| raise NotImplementedError |
|
|
|
|
| @torch.inference_mode() |
| def diffuse( |
| model: AnonymousEncoder, |
| input_ids: Tensor, |
| editable_mask: Tensor, |
| *, |
| steps: int, |
| mask_id: int, |
| pad_id: int, |
| temperature: float = 1.0, |
| top_k: int | None = None, |
| forbidden_token_ids: tuple[int, ...] = (), |
| generator: torch.Generator | None = None, |
| ) -> DiffusionOutput: |
| """Run all reverse steps, returning the final sequence and its trajectory.""" |
| raise NotImplementedError |
|
|
