File size: 1,982 Bytes
d441d06 b21f2e1 d441d06 | 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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """Student scaffold for absorbing-mask diffusion sampling."""
from __future__ import annotations
from dataclasses import dataclass
import torch
from torch import Tensor
# Prerequisite: copy the model.py that you completed in the Encoder lab into
# this directory. Its solution is intentionally not included in this lab.
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
|