diffuusio & h-net kielimalli

diffusion_model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# Notice that this one isn't a working piece. I have another coming up that is much more promising.

# ── H-Net context encoder ─────────────────────────────────────────────────────────
# Pulled from hnet.py — only the components needed for context compression.
# Full H-Net (with main network + decoder) is not used here; we only want the
# encoder + dynamic chunking to compress the raw-byte context embeddings into
# a shorter sequence of semantically richer chunk representations.

class _RMSNorm(nn.Module):
    def __init__(self, d: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        return self.weight * x / x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()


class _RoutingModule(nn.Module):
    """Cosine-similarity boundary predictor (H-Net eq. 4)."""
    def __init__(self, d: int):
        super().__init__()
        self.W_q = nn.Linear(d, d, bias=False)
        self.W_k = nn.Linear(d, d, bias=False)

    def forward(self, x):
        q = self.W_q(x)
        k = self.W_k(x)
        k_prev = torch.cat([torch.zeros_like(k[:, :1]), k[:, :-1]], dim=1)
        cos = F.cosine_similarity(q, k_prev, dim=-1)
        p = 0.5 * (1.0 - cos)
        p = torch.cat([torch.ones_like(p[:, :1]), p[:, 1:]], dim=1)
        b = (p >= 0.5).float()
        return p, b


class _SmoothingModule(nn.Module):
    """
    EMA-based differentiable upsampling (H-Net eq. 5):
        z_bar_t = P_t * z_t + (1 - P_t) * z_bar_{t-1}

    Implemented as a parallel prefix scan so the entire sequence is processed
    in O(log L) sequential steps instead of O(L), which matters at CONTEXT_LEN=64.

    The recurrence z_bar_t = a_t * z_bar_{t-1} + b_t has:
        a_t = (1 - P_t)   (decay / carry-over coefficient)
        b_t = P_t * z_t   (input term)

    Parallel scan identity: (a1, b1) ∘ (a2, b2) = (a1*a2, a1*b2 + b1)
    """
    def forward(self, z: torch.Tensor, P: torch.Tensor) -> torch.Tensor:
        B, L, D = z.shape
        a = (1.0 - P).unsqueeze(-1).expand(B, L, D).contiguous()  # (B, L, D)
        b = P.unsqueeze(-1) * z                                     # (B, L, D)

        # Parallel prefix scan — double stride each pass, O(log L) steps
        stride = 1
        while stride < L:
            a_prev = torch.cat([a[:, :stride], a[:, :-stride]], dim=1)
            b_prev = torch.cat([b[:, :stride], b[:, :-stride]], dim=1)
            new_a = a.clone()
            new_b = b.clone()
            new_a[:, stride:] = a[:, stride:] * a_prev[:, stride:]
            new_b[:, stride:] = a[:, stride:] * b_prev[:, stride:] + b[:, stride:]
            a, b = new_a, new_b
            stride *= 2

        return b  # prefix-scan result = z_bar at each position


class _SWAttention(nn.Module):
    """Causal sliding-window self-attention."""
    def __init__(self, d: int, n_heads: int, window: int = 1024):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d // n_heads
        self.window = window
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.out = nn.Linear(d, d, bias=False)

    def forward(self, x):
        B, L, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=-1)
        q = q.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        mask = torch.triu(torch.ones(L, L, device=x.device, dtype=torch.bool), diagonal=1)
        dist = torch.arange(L, device=x.device)
        mask = mask | ((dist.unsqueeze(0) - dist.unsqueeze(1)) < -self.window)
        scale = self.head_dim ** -0.5
        attn = (q @ k.transpose(-2, -1)) * scale
        attn = attn.masked_fill(mask[None, None], float("-inf"))
        attn = F.softmax(attn, dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, L, D)
        return self.out(out)


class _EncBlock(nn.Module):
    def __init__(self, d: int, n_heads: int, window: int = 1024):
        super().__init__()
        self.norm1 = _RMSNorm(d)
        self.attn  = _SWAttention(d, n_heads, window)
        self.norm2 = _RMSNorm(d)
        d_ff = (int(d * 8 / 3) + 63) // 64 * 64
        self.gate = nn.Linear(d, d_ff, bias=False)
        self.up   = nn.Linear(d, d_ff, bias=False)
        self.down = nn.Linear(d_ff, d, bias=False)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.down(F.silu(self.gate(self.norm2(x))) * self.up(self.norm2(x)))
        return x


class HNetContextEncoder(nn.Module):
    """
    H-Net encoder + dynamic chunking used as a context compressor.

    Takes byte-level embeddings (B, L, d_in) and produces compressed chunk
    representations (B, L', d_out) where L' ≈ L / N.

    This is the encoder half of H-Net only — no main network, no decoder,
    no LM head. Its job is purely to build richer, shorter context for the
    diffusion denoiser's KV cache.

    Args:
        d_in:     input embedding dimension (must match denoiser d_model)
        d_out:    output chunk dimension (can differ; projected if needed)
        n_layers: number of encoder transformer layers (paper uses 4)
        n_heads:  attention heads
        N:        target compression ratio (paper: ~6 bytes/chunk for English)
        window:   sliding window size for encoder attention
    """
    def __init__(
        self,
        d_in:     int,
        d_out:    int,
        n_layers: int = 4,
        n_heads:  int = 8,
        N:        float = 4.0,
        window:   int = 1024,
    ):
        super().__init__()
        self.N = N
        self.d_out = d_out

        self.layers = nn.ModuleList([_EncBlock(d_in, n_heads, window) for _ in range(n_layers)])
        self.norm   = _RMSNorm(d_in)
        self.router = _RoutingModule(d_in)
        self.smoother = _SmoothingModule()

        # Project to d_out if dimensions differ
        self.proj = nn.Linear(d_in, d_out, bias=False) if d_in != d_out else nn.Identity()

        # Ratio loss weight (same as hnet.py alpha)
        self.alpha = 0.03

    def encode(self, x: torch.Tensor):
        """
        Args:
            x: (B, L, d_in) byte-level embeddings
        Returns:
            chunks:      (B, L', d_out) compressed representations, L' ≈ L/N
            ratio_loss:  scalar — steers compression toward target ratio
        """
        # Run encoder layers
        h = x
        for layer in self.layers:
            h = layer(h)
        x_hat = self.norm(h)                          # (B, L, d_in)

        # Dynamic chunking: predict boundaries, select boundary positions
        p, b = self.router(x_hat)                     # (B, L) each

        # Compress: keep only boundary-marked vectors
        B, L, D = x_hat.shape
        b_bool = b.bool()
        lengths = b_bool.sum(dim=1)                   # (B,)
        max_len = max(lengths.max().item(), 1)

        x_comp = torch.zeros(B, max_len, D, device=x.device, dtype=x.dtype)
        for i in range(B):
            sel = b_bool[i].nonzero(as_tuple=False).squeeze(-1)
            n = sel.size(0)
            if n > 0:
                x_comp[i, :n] = x_hat[i, sel]

        # Project to output dimension
        chunks = self.proj(x_comp)                    # (B, L', d_out)

        # Ratio loss (H-Net eq. 10)
        F_val = b.float().mean(dim=1)                 # fraction selected (non-diff)
        G_val = p.mean(dim=1)                         # mean boundary prob (diff)
        N = self.N
        ratio_loss = (N / (N - 1) * ((N - 1) * F_val * G_val + (1 - F_val) * (1 - G_val))).mean()

        return chunks, ratio_loss, lengths

    def forward(self, x: torch.Tensor):
        """Convenience wrapper — returns only chunks and ratio_loss."""
        chunks, ratio_loss, _ = self.encode(x)
        return chunks, ratio_loss


CONTEXT_LEN = 64

# ── Noise schedule ───────────────────────────────────────────────────────────────
def cosine_beta_schedule(T: int, s: float = 0.008):
    steps = torch.arange(T + 1, dtype=torch.float64)
    alphas_cumprod = torch.cos(((steps / T) + s) / (1 + s) * math.pi / 2) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0.0001, 0.9999).float()


class GaussianDiffusion:
    def __init__(self, T: int = 200):
        self.T = T
        betas      = cosine_beta_schedule(T)
        alphas     = 1.0 - betas
        alphas_bar = torch.cumprod(alphas, dim=0)
        self.betas          = betas
        self.alphas_bar     = alphas_bar
        self.sqrt_ab        = alphas_bar.sqrt()
        self.sqrt_one_minus = (1 - alphas_bar).sqrt()

    def _to(self, device):
        self.betas          = self.betas.to(device)
        self.alphas_bar     = self.alphas_bar.to(device)
        self.sqrt_ab        = self.sqrt_ab.to(device)
        self.sqrt_one_minus = self.sqrt_one_minus.to(device)
        return self

    def q_sample(self, x0: torch.Tensor, t: torch.Tensor):
        self._to(x0.device)
        noise = torch.randn_like(x0)
        sab = self.sqrt_ab[t][:, None, None]
        som = self.sqrt_one_minus[t][:, None, None]
        return sab * x0 + som * noise, noise

    @torch.no_grad()
    def ddim_sample_one(self, model, context: torch.Tensor,
                        steps: int = 10, eta: float = 0.0) -> torch.Tensor:
        """
        DDIM inference KV-cachella.
        Kontekstin K/V lasketaan kerran, target-query per askel.
        """
        self._to(context.device)
        B, ctx_len, d = context.shape

        # Laske kontekstin KV-cache kerran — pysyy vakiona koko DDIM-loopin ajan
        kv_cache = model.compute_kv_cache(context)

        indices = torch.linspace(0, self.T - 1, steps).long().flip(0)
        x = torch.randn(B, 1, d, device=context.device)

        for i, t in enumerate(indices):
            t_val    = t.item()
            t_tensor = torch.full((B,), t_val, device=context.device, dtype=torch.long)

            # Forward vain target-positiolle, konteksti KV-cachesta
            x0_pred = model.forward_with_cache(x, kv_cache, t_tensor)

            ab_t = self.alphas_bar[t_val]
            eps_pred = (x - ab_t.sqrt() * x0_pred) / (1 - ab_t).sqrt().clamp(min=1e-8)

            if i == len(indices) - 1:
                x = x0_pred
            else:
                ab_tm1 = self.alphas_bar[indices[i + 1].item()]
                sigma  = eta * ((1 - ab_tm1) / (1 - ab_t) * (1 - ab_t / ab_tm1)).sqrt()
                x = (ab_tm1.sqrt() * x0_pred
                     + (1 - ab_tm1 - sigma**2).clamp(min=0).sqrt() * eps_pred
                     + sigma * torch.randn_like(x))
        return x


# ── Timestep embedding ────────────────────────────────────────────────────────────
class TimestepEmbedding(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.SiLU(),
            nn.Linear(d_model * 4, d_model),
        )
        half = d_model // 2
        freqs = torch.exp(-math.log(10000) * torch.arange(half) / (half - 1))
        self.register_buffer("freqs", freqs)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        emb = t[:, None].float() * self.freqs[None]
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
        return self.proj(emb)


# ── Transformer-kerros KV-cachella ───────────────────────────────────────────────
class CachedAttentionLayer(nn.Module):
    """
    Pre-norm transformer-kerros joka tukee kahta forward-moodia:

    1. train / full forward:
       forward(seq)  →  (B, L, d)
       Normaali full-sequence forward koulutukseen.

    2. cache-moodi inference:
       compute_kv(context)        →  (K, V) tuple
       forward_cached(x_t, K, V)  →  (B, 1, d)
       x_t on vain target-positio, K/V on kontekstin esilasketut avaimet/arvot.
       Attention: query = x_t,  key/value = cat(context_kv, x_t_kv)
    """
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model  = d_model
        self.n_heads  = n_heads
        self.d_head   = d_model // n_heads

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        # Yhdistetty QKV-projektio — yksi matriisikertolasku kolmen sijaan
        self.qkv    = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out    = nn.Linear(d_model, d_model, bias=False)

        # Erilliset Q/KV-projektiot cache-moodiin
        self.q_proj  = nn.Linear(d_model, d_model, bias=False)
        self.kv_proj = nn.Linear(d_model, 2 * d_model, bias=False)

        self.ff = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model),
            nn.Dropout(dropout),
        )
        self.attn_drop = dropout

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        """(B, L, d) → (B, n_heads, L, d_head)"""
        B, L, _ = x.shape
        return x.view(B, L, self.n_heads, self.d_head).transpose(1, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Full forward — koulutuksessa."""
        B, L, _ = x.shape
        h = self.norm1(x)
        qkv = self.qkv(h).chunk(3, dim=-1)
        q, k, v = [self._split_heads(t) for t in qkv]

        # Flash attention jos saatavilla (PyTorch 2.0+)
        attn_out = F.scaled_dot_product_attention(
            q, k, v, dropout_p=self.attn_drop if self.training else 0.0
        )
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, L, self.d_model)
        x = x + self.out(attn_out)
        x = x + self.ff(self.norm2(x))
        return x

    def compute_kv(self, context: torch.Tensor):
        """
        Esilaskee kontekstin K ja V inference-cacheen.
        context: (B, ctx_len, d)
        Palauttaa: (K, V) molemmat (B, n_heads, ctx_len, d_head)
        """
        h  = self.norm1(context)
        kv = self.kv_proj(h)
        k, v = kv.chunk(2, dim=-1)
        return self._split_heads(k), self._split_heads(v)

    def forward_cached(self, x_t: torch.Tensor, ctx_k, ctx_v) -> torch.Tensor:
        """
        Forward vain target-positiolle käyttäen cachettuja K/V:tä.
        x_t: (B, 1, d)
        ctx_k, ctx_v: (B, n_heads, ctx_len, d_head)
        """
        B = x_t.shape[0]
        h = self.norm1(x_t)

        # Query vain target-positiosta
        q = self._split_heads(self.q_proj(h))   # (B, n_heads, 1, d_head)

        # Key/Value: konteksti + target yhdistettynä
        kv     = self.kv_proj(h)
        tk, tv = kv.chunk(2, dim=-1)
        tk     = self._split_heads(tk)           # (B, n_heads, 1, d_head)
        tv     = self._split_heads(tv)

        k = torch.cat([ctx_k, tk], dim=2)        # (B, n_heads, ctx_len+1, d_head)
        v = torch.cat([ctx_v, tv], dim=2)

        attn_out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, 1, self.d_model)

        x_t = x_t + self.out(attn_out)
        x_t = x_t + self.ff(self.norm2(x_t))
        return x_t


# ── Denoiser ──────────────────────────────────────────────────────────────────────
class ARDiffusionDenoiser(nn.Module):
    def __init__(
        self,
        vocab_size:   int,
        d_model:      int   = 256,
        n_heads:      int   = 8,
        n_layers:     int   = 6,
        dropout:      float = 0.1,
        hnet_encoder: "HNetContextEncoder | None" = None,
    ):
        super().__init__()
        self.d_model   = d_model
        self.n_layers  = n_layers
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.emb_scale = nn.Parameter(torch.ones(1))
        self.time_emb  = TimestepEmbedding(d_model)

        # H-Net encoder: compresses context before KV cache is built.
        # When provided, context length shrinks from CONTEXT_LEN to ~CONTEXT_LEN/N,
        # so pos_emb must cover the full uncompressed length (worst case = no compression).
        # We use a learned pos_emb over CONTEXT_LEN+1 slots regardless; for compressed
        # contexts we interpolate positions to [0, CONTEXT_LEN-1] so the denoiser always
        # sees familiar position indices no matter the compression ratio.
        self.hnet_encoder = hnet_encoder
        self.pos_emb   = nn.Embedding(CONTEXT_LEN + 1, d_model)

        self.layers   = nn.ModuleList([
            CachedAttentionLayer(d_model, n_heads, dropout)
            for _ in range(n_layers)
        ])
        self.norm     = nn.LayerNorm(d_model)
        self.out_proj = nn.Linear(d_model, d_model)

        # Esilaskettu pos-embedding — vakio, ei muutu
        pos = torch.arange(CONTEXT_LEN + 1)
        self.register_buffer("_pos_ids", pos)

    def get_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor:
        emb = self.embedding(token_ids)
        return F.normalize(emb, dim=-1) * self.emb_scale

    def _compress_context(self, context: torch.Tensor):
        """
        If hnet_encoder is set, compress context embeddings to chunk representations.
        Returns (compressed_context, ratio_loss).

        The compressed context has shape (B, L', d) where L' ≤ CONTEXT_LEN.
        Positional embeddings are re-interpolated to evenly span [0, CONTEXT_LEN-1]
        regardless of L', so the denoiser always sees familiar position ranges.

        Without hnet_encoder this is a no-op (ratio_loss = 0).
        """
        if self.hnet_encoder is None:
            return context, torch.tensor(0.0, device=context.device)

        compressed, ratio_loss = self.hnet_encoder(context)  # (B, L', d)
        return compressed, ratio_loss

    def _context_pos_emb(self, L_ctx: int, device) -> torch.Tensor:
        """
        Positional embeddings for a context of length L_ctx.
        Indices are evenly spread across [0, CONTEXT_LEN-1] so that regardless
        of whether the context was compressed or not, position semantics stay
        roughly aligned with what the model learned.
        """
        if L_ctx == CONTEXT_LEN:
            return self.pos_emb(self._pos_ids[:CONTEXT_LEN])   # (L_ctx, d)
        # Interpolate: map L' positions → CONTEXT_LEN slots (float → round)
        idx = torch.linspace(0, CONTEXT_LEN - 1, L_ctx, device=device).long()
        return self.pos_emb(idx)                                # (L_ctx, d)

    # ── Koulutus: full forward ────────────────────────────────────────────────────
    def forward(self, x_t: torch.Tensor, context: torch.Tensor, t: torch.Tensor):
        """
        x_t:     (B, 1, d)
        context: (B, CONTEXT_LEN, d)  — raw byte embeddings
        t:       (B,)

        Returns: (B, 1, d) denoised prediction, ratio_loss scalar
        """
        # Compress context through H-Net encoder (no-op if encoder is None)
        ctx, ratio_loss = self._compress_context(context)      # (B, L', d)
        B, L_ctx, _ = ctx.shape

        pos_ctx = self._context_pos_emb(L_ctx, ctx.device)     # (L', d)
        ctx = ctx + pos_ctx.unsqueeze(0)

        # Target token gets the last positional slot + timestep embedding
        x_t_pos = x_t + self.pos_emb(self._pos_ids[CONTEXT_LEN:CONTEXT_LEN+1]).unsqueeze(0)
        x_t_pos = x_t_pos + self.time_emb(t).unsqueeze(1)

        seq = torch.cat([ctx, x_t_pos], dim=1)                 # (B, L'+1, d)

        for layer in self.layers:
            seq = layer(seq)

        out = self.out_proj(self.norm(seq[:, -1:]))             # (B, 1, d)
        return out, ratio_loss

    # ── Inference: KV-cache ───────────────────────────────────────────────────────
    def compute_kv_cache(self, context: torch.Tensor) -> list:
        """
        Esilaskee kaikille kerroksille kontekstin K ja V.
        context: (B, CONTEXT_LEN, d)  — raw byte embeddings

        With hnet_encoder: context is compressed to (B, L', d) first.
        L' ≈ CONTEXT_LEN / N — shorter sequence, richer representations,
        faster attention during DDIM loop.
        """
        # Compress (or pass through unchanged if no encoder)
        ctx, _ = self._compress_context(context)               # (B, L', d)
        B, L_ctx, _ = ctx.shape

        pos_ctx = self._context_pos_emb(L_ctx, ctx.device)
        ctx = ctx + pos_ctx.unsqueeze(0)

        cache = []
        h = ctx
        for layer in self.layers:
            k, v = layer.compute_kv(h)
            cache.append((k, v))
            h = layer(h)

        return cache

    @torch.no_grad()
    def forward_with_cache(self, x_t: torch.Tensor, kv_cache: list, t: torch.Tensor) -> torch.Tensor:
        """
        Yksi denoising-askel KV-cachella.
        x_t:      (B, 1, d)
        kv_cache: lista (K, V) per kerros — konteksti pysyy vakiona
        t:        (B,)
        """
        # Target-position embedding (same slot as during training)
        h = x_t + self.pos_emb(self._pos_ids[CONTEXT_LEN:CONTEXT_LEN+1]).unsqueeze(0)
        h = h + self.time_emb(t).unsqueeze(1)

        for layer, (ctx_k, ctx_v) in zip(self.layers, kv_cache):
            h = layer.forward_cached(h, ctx_k, ctx_v)

        return self.out_proj(self.norm(h))


# ── Training loss ─────────────────────────────────────────────────────────────────
_step_counter = [0]

def diffusion_loss(
    model:          ARDiffusionDenoiser,
    diffusion:      GaussianDiffusion,
    token_ids:      torch.Tensor,
    rounding_weight: float = 0.1,
    rounding_every:  int   = 4,
) -> torch.Tensor:
    B, L   = token_ids.shape
    device = token_ids.device

    x0_full = model.get_embeddings(token_ids)   # (B, L, d)

    pos     = torch.randint(CONTEXT_LEN, L, (B,), device=device)
    ctx_idx = pos[:, None] + torch.arange(-CONTEXT_LEN, 0, device=device)[None, :]
    b_idx   = torch.arange(B, device=device)

    context = x0_full[b_idx[:, None], ctx_idx]          # (B, CTX, d)
    target  = x0_full[b_idx, pos].unsqueeze(1)           # (B, 1, d)

    t      = torch.randint(0, diffusion.T, (B,), device=device)
    x_t, _ = diffusion.q_sample(target.detach(), t)

    # model.forward now returns (prediction, ratio_loss).
    # context is NOT detached here so gradients flow into hnet_encoder.
    x0_pred, ratio_loss = model(x_t, context, t)         # (B, 1, d), scalar

    loss = F.mse_loss(x0_pred, target.detach())

    # H-Net ratio loss: steers the encoder toward the target compression ratio.
    # Only non-zero when hnet_encoder is attached.
    if model.hnet_encoder is not None:
        loss = loss + model.hnet_encoder.alpha * ratio_loss

    _step_counter[0] += 1
    if rounding_weight > 0 and (_step_counter[0] % rounding_every == 0):
        W          = F.normalize(model.embedding.weight, dim=-1) * model.emb_scale
        logits     = x0_pred[:, 0] @ W.T / 0.1
        target_ids = token_ids[b_idx, pos]
        loss       = loss + rounding_weight * F.cross_entropy(logits, target_ids)

    return loss


# ── Decode ────────────────────────────────────────────────────────────────────────
@torch.no_grad()
def decode(model: ARDiffusionDenoiser, x0: torch.Tensor) -> torch.Tensor:
    W = F.normalize(model.embedding.weight, dim=-1) * model.emb_scale
    return (x0 @ W.T).argmax(-1)


# ── Masked diffusion loss ─────────────────────────────────────────────────────────
def masked_diffusion_loss(
    model:           ARDiffusionDenoiser,
    diffusion:       GaussianDiffusion,
    token_ids:       torch.Tensor,
    mask_ratio:      float = 0.5,
    rounding_weight: float = 0.5,
) -> torch.Tensor:
    """
    Infilling training objective.

    Calls backward() per sample so only one sample's computation graph
    lives in VRAM at a time. Returns a plain scalar tensor for logging;
    gradients are already accumulated into model.parameters() on return.
    """
    B, L   = token_ids.shape
    device = token_ids.device

    total_loss_val = 0.0
    n_valid        = 0

    for b in range(B):
        ids_b = token_ids[b]
        x0_b  = model.get_embeddings(ids_b.unsqueeze(0))[0]    # (L, d)

        mask = torch.rand(L, device=device) < mask_ratio
        mask[torch.randint(L, (1,))] = False                    # ≥1 anchor
        mask[torch.randint(L, (1,))] = True                     # ≥1 target

        target_positions = mask.nonzero(as_tuple=False).squeeze(-1)
        if target_positions.numel() == 0:
            continue
        if target_positions.numel() > 16:
            perm = torch.randperm(target_positions.numel(), device=device)[:16]
            target_positions = target_positions[perm]

        t_val       = torch.randint(0, diffusion.T, (1,), device=device)
        x_noised, _ = diffusion.q_sample(x0_b.unsqueeze(0), t_val)
        x_noised    = x_noised[0]                               # (L, d)

        mask_f  = mask.float().unsqueeze(-1)
        x_input = mask_f * x_noised + (1 - mask_f) * x0_b     # (L, d)

        ctx, ratio_loss = model._compress_context(x_input.unsqueeze(0))
        ctx = ctx + model._context_pos_emb(ctx.shape[1], device).unsqueeze(0)

        n         = target_positions.numel()
        t_emb_b   = model.time_emb(t_val)
        pos_last  = model.pos_emb(model._pos_ids[CONTEXT_LEN:CONTEXT_LEN+1]).unsqueeze(0)

        h   = x_noised[target_positions].unsqueeze(1) + pos_last + t_emb_b.expand(n, -1).unsqueeze(1)
        seq = torch.cat([ctx.expand(n, -1, -1), h], dim=1)
        for layer in model.layers:
            seq = layer(seq)
        x0_pred = model.out_proj(model.norm(seq[:, -1:]))       # (n, 1, d)

        x0_targets = x0_b[target_positions].unsqueeze(1).detach()
        mse        = F.mse_loss(x0_pred, x0_targets)
        W          = F.normalize(model.embedding.weight, dim=-1) * model.emb_scale
        round_loss = F.cross_entropy(x0_pred[:, 0] @ W.T / 0.1, ids_b[target_positions])

        sample_loss = (mse + rounding_weight * round_loss) / B
        if model.hnet_encoder is not None:
            sample_loss = sample_loss + model.hnet_encoder.alpha * ratio_loss / B

        # Backward now — frees this sample's graph immediately
        sample_loss.backward()
        total_loss_val += sample_loss.item()
        n_valid        += 1

        del seq, x0_pred, ctx, x_input, x_noised, x0_b, h, sample_loss
        torch.cuda.empty_cache()

    return torch.tensor(total_loss_val, device=device)


# ── Infilling inference ───────────────────────────────────────────────────────────
@torch.no_grad()
def infill_sample(
    text_with_gaps: str,
    model:          ARDiffusionDenoiser,
    diffusion:      GaussianDiffusion,
    gap_char:       str   = "_",
    ddim_steps:     int   = 20,
    eta:            float = 0.0,
    device:         str   = "cpu",
) -> str:
    """
    Fill gaps in a text template.

    Mark positions to be generated with the gap_char (default '_').
    Each gap character = one byte to generate.

    Example:
        infill_sample("talo__ kadulla", model, diffusion)
        # might return "talossa kadulla"

    The model sees the full template (anchors clean, gaps noised) compressed
    through H-Net, giving it bidirectional context to fill each gap.

    For best results, keep the total length ≤ CONTEXT_LEN bytes and use
    a model trained with masked_diffusion_loss.
    """
    model.eval()
    dev = torch.device(device)
    model.to(dev)

    # Encode template: anchor bytes get their real embeddings, gaps start as noise
    raw_bytes  = text_with_gaps.encode("utf-8")
    L          = len(raw_bytes)
    d          = model.d_model

    # Identify gap positions and anchor positions
    gap_byte   = ord(gap_char)
    is_gap     = torch.tensor([b == gap_byte for b in raw_bytes], device=dev)  # (L,)
    anchor_ids = torch.tensor(
        [b if b != gap_byte else 0 for b in raw_bytes],
        dtype=torch.long, device=dev
    )                                                       # (L,) — gaps get dummy id 0

    anchor_embs = model.get_embeddings(anchor_ids)          # (L, d)

    # Pad or truncate to CONTEXT_LEN
    if L < CONTEXT_LEN:
        pad_emb = torch.zeros(CONTEXT_LEN - L, d, device=dev)
        anchor_embs = torch.cat([pad_emb, anchor_embs], dim=0)
        pad_gap = torch.zeros(CONTEXT_LEN - L, dtype=torch.bool, device=dev)
        is_gap  = torch.cat([pad_gap, is_gap], dim=0)
    else:
        anchor_embs = anchor_embs[-CONTEXT_LEN:]
        is_gap      = is_gap[-CONTEXT_LEN:]

    anchor_embs = anchor_embs.unsqueeze(0)                  # (1, CONTEXT_LEN, d)
    is_gap      = is_gap.unsqueeze(0)                       # (1, CONTEXT_LEN)

    # Initialise gap positions with pure noise
    x = anchor_embs.clone()
    x[is_gap] = torch.randn(is_gap.sum(), d, device=dev)

    # DDIM loop over all gap positions simultaneously
    indices = torch.linspace(0, diffusion.T - 1, ddim_steps).long().flip(0)
    diffusion._to(dev)

    for i, t_val in enumerate(indices):
        t_val   = t_val.item()
        t_tensor = torch.full((1,), t_val, device=dev, dtype=torch.long)

        # Compress the full window (anchors clean, gaps at current noise level)
        ctx, _ = model._compress_context(x)                # (1, L', d)
        pos_ctx = model._context_pos_emb(ctx.shape[1], dev)
        ctx     = ctx + pos_ctx.unsqueeze(0)

        # Denoise each gap position
        x0_preds = torch.zeros_like(x)                     # (1, CONTEXT_LEN, d)
        gap_positions = is_gap[0].nonzero(as_tuple=False).squeeze(-1)

        for pos in gap_positions:
            x_t  = x[0, pos].unsqueeze(0).unsqueeze(0)    # (1, 1, d)
            h    = x_t + model.pos_emb(model._pos_ids[CONTEXT_LEN:CONTEXT_LEN+1]).unsqueeze(0)
            h    = h + model.time_emb(t_tensor).unsqueeze(1)
            seq  = torch.cat([ctx, h], dim=1)
            for layer in model.layers:
                seq = layer(seq)
            x0_pred = model.out_proj(model.norm(seq[:, -1:]))  # (1, 1, d)
            x0_preds[0, pos] = x0_pred[0, 0]

        # DDIM step for all gap positions
        ab_t = diffusion.alphas_bar[t_val]
        if i == len(indices) - 1:
            x[is_gap] = x0_preds[is_gap]
        else:
            ab_tm1 = diffusion.alphas_bar[indices[i + 1].item()]
            eps    = (x - ab_t.sqrt() * x0_preds) / (1 - ab_t).sqrt().clamp(min=1e-8)
            sigma  = eta * ((1 - ab_tm1) / (1 - ab_t) * (1 - ab_t / ab_tm1)).sqrt()
            x_next = (ab_tm1.sqrt() * x0_preds
                      + (1 - ab_tm1 - sigma**2).clamp(min=0).sqrt() * eps
                      + sigma * torch.randn_like(x))
            # Only update gap positions; keep anchors exactly clean
            x[is_gap] = x_next[is_gap]

    # Decode all positions
    ids = decode(model, x)                                  # (1, CONTEXT_LEN)
    # Extract only the filled-in gap region (remove left padding if any)
    offset = max(0, CONTEXT_LEN - L)
    result_ids = ids[0, offset:].cpu().tolist()

    # Reconstruct: anchors keep original bytes, gaps get decoded bytes
    orig_gap_positions = [j for j, b in enumerate(raw_bytes) if b == gap_byte]
    result = bytearray(raw_bytes)
    filled_gap_idx = 0
    gap_positions_in_ctx = is_gap[0, offset:].nonzero(as_tuple=False).squeeze(-1).tolist()
    for ctx_pos, orig_pos in zip(gap_positions_in_ctx, orig_gap_positions):
        result[orig_pos] = result_ids[ctx_pos] & 0xFF

    return bytes(result).decode("utf-8", errors="replace")


# ── AR sampling ───────────────────────────────────────────────────────────────────
@torch.no_grad()
def ar_sample(
    model:       ARDiffusionDenoiser,
    diffusion:   GaussianDiffusion,
    n_steps:     int,
    device,
    seed_tokens: torch.Tensor = None,
    ddim_steps:  int   = 10,
    eta:         float = 0.0,
) -> torch.Tensor:
    d = model.d_model

    if seed_tokens is not None:
        context = model.get_embeddings(seed_tokens.to(device))
    else:
        context = torch.zeros(1, CONTEXT_LEN, d, device=device)

    generated = []
    for _ in range(n_steps):
        # ddim_sample_one calls compute_kv_cache internally.
        # If hnet_encoder is set, context is compressed there before building
        # the KV cache — so each DDIM step attends to ~CONTEXT_LEN/N keys
        # instead of CONTEXT_LEN keys.
        x0 = diffusion.ddim_sample_one(model, context, steps=ddim_steps, eta=eta)
        generated.append(x0)
        # Rolling window stays in raw byte-embedding space; re-compressed each step.
        context = torch.cat([context[:, 1:], x0], dim=1)

    return torch.cat(generated, dim=1)


# ── Checkpoint save / load ────────────────────────────────────────────────────────
import os

def save_checkpoint(path: str, epoch: int, model, optimizer, scheduler, loss: float):
    torch.save({
        "epoch":          epoch,
        "model":          model.state_dict(),
        "optimizer":      optimizer.state_dict(),
        "scheduler":      scheduler.state_dict(),
        "loss":           loss,
    }, path)
    print(f"  ✓ checkpoint saved → {path}")


def load_checkpoint(path: str, model, optimizer, scheduler):
    ckpt = torch.load(path, weights_only=True)
    model.load_state_dict(ckpt["model"])
    optimizer.load_state_dict(ckpt["optimizer"])
    try:
        scheduler.load_state_dict(ckpt["scheduler"])
    except (KeyError, ValueError):
        # Scheduler type changed (e.g. CosineAnnealingLR → LambdaLR) — skip it
        # and let the new schedule start from epoch 0. Weights and optimizer
        # momentum are preserved, which is what actually matters for resuming.
        print("  ! scheduler state skipped (type changed — starting LR schedule fresh)")
    print(f"  ✓ resumed from {path}  (epoch {ckpt['epoch']}, loss {ckpt['loss']:.4f})")
    return ckpt["epoch"] + 1


# ── Completion API ───────────────────────────────────────────────────────────────
def complete(
    prompt:      str,
    model:       "ARDiffusionDenoiser",
    diffusion:   "GaussianDiffusion",
    n_tokens:    int   = 80,
    ddim_steps:  int   = 10,
    eta:         float = 0.0,
    device:      str   = "cpu",
) -> str:
    """
    Complete a text prompt using a trained model.

    Usage after importing:
        from diffusion_model8 import complete, load_model
        model, diffusion = load_model("checkpoints/ckpt_epoch0099.pt")
        print(complete("Vaka vanha Väinämöinen", model, diffusion))

    Args:
        prompt:     seed text; last CONTEXT_LEN bytes are used as context
        model:      trained ARDiffusionDenoiser (call load_model() to get one)
        diffusion:  GaussianDiffusion instance
        n_tokens:   number of bytes to generate
        ddim_steps: DDIM denoising steps per token (10 is fast, 50 is higher quality)
        eta:        DDIM stochasticity — 0.0 = deterministic, 1.0 = full noise
        device:     "cuda" or "cpu"

    Returns:
        prompt + generated continuation as a string
    """
    model.eval()
    dev = torch.device(device)
    model.to(dev)

    # Encode prompt bytes; pad with zeros on the left if shorter than CONTEXT_LEN
    prompt_bytes = prompt.encode("utf-8")
    raw = torch.tensor(list(prompt_bytes), dtype=torch.long, device=dev)

    if raw.shape[0] >= CONTEXT_LEN:
        seed = raw[-CONTEXT_LEN:]                              # take last CONTEXT_LEN bytes
    else:
        pad  = torch.zeros(CONTEXT_LEN - raw.shape[0], dtype=torch.long, device=dev)
        seed = torch.cat([pad, raw])

    # Build initial context as normalised embeddings, shape (1, CONTEXT_LEN, d)
    context = model.get_embeddings(seed.unsqueeze(0))          # (1, CONTEXT_LEN, d)

    generated_ids = []
    with torch.no_grad():
        for _ in range(n_tokens):
            x0 = diffusion.ddim_sample_one(model, context, steps=ddim_steps, eta=eta)
            ids = decode(model, x0)                            # (1, 1)
            generated_ids.append(ids[0, 0].item())
            context = torch.cat([context[:, 1:], x0], dim=1)  # slide window

    generated_bytes = bytes(generated_ids)
    return prompt + generated_bytes.decode("utf-8", errors="replace")


def load_model(
    checkpoint:  str,
    d_model:     int   = 384,
    n_layers:    int   = 6,
    n_heads:     int   = 8,
    vocab_size:  int   = 256,
    hnet_N:      float = 4.0,
    device:      str   = "cpu",
) -> tuple:
    """
    Rebuild the model architecture and load weights from a checkpoint.

    Returns (model, diffusion) ready to pass to complete().
    """
    hnet_encoder = HNetContextEncoder(
        d_in=d_model, d_out=d_model,
        n_layers=4, n_heads=n_heads,
        N=hnet_N, window=CONTEXT_LEN,
    )
    model = ARDiffusionDenoiser(
        vocab_size=vocab_size,
        d_model=d_model,
        n_heads=n_heads,
        n_layers=n_layers,
        hnet_encoder=hnet_encoder,
    )
    ckpt = torch.load(checkpoint, map_location=device, weights_only=True)
    model.load_state_dict(ckpt["model"])
    model.to(device).eval()

    diffusion = GaussianDiffusion(T=200)
    epoch = ckpt.get("epoch", "?")
    loss  = ckpt.get("loss",  float("nan"))
    print(f"Loaded checkpoint: {checkpoint}  (epoch {epoch}, loss {loss:.4f})")
    return model, diffusion


# ── Main ──────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    device  = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    VOCAB   = 256
    D_MODEL = 384
    BATCH   = 48           # reduced from 96 — masked loss needs headroom alongside AR
    SEQ_LEN = CONTEXT_LEN + 32
    SAVE_DIR      = "checkpoints"
    SAVE_EVERY    = 5
    RESUME_LATEST = True   # set False to always start from scratch

    os.makedirs(SAVE_DIR, exist_ok=True)

    diffusion = GaussianDiffusion(T=200)

    hnet_encoder = HNetContextEncoder(
        d_in=D_MODEL, d_out=D_MODEL,
        n_layers=4, n_heads=8,
        N=4.0,
        window=CONTEXT_LEN,
    )

    model = ARDiffusionDenoiser(
        vocab_size=VOCAB,
        d_model=D_MODEL,
        n_heads=8,
        n_layers=6,
        hnet_encoder=hnet_encoder,
    ).to(device)

    total_params = sum(p.numel() for p in model.parameters())
    enc_params   = sum(p.numel() for p in hnet_encoder.parameters())
    print(f"HNet encoder    : {enc_params:,}")
    print(f"Denoiser        : {total_params - enc_params:,}")
    print(f"Total           : {total_params:,}")
    print(f"Context         : {CONTEXT_LEN} bytes → ~{CONTEXT_LEN // 4} chunks")
    print(f"Device          : {device}")
    print()

    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

    # Warmup-Stable-Decay schedule:
    #   - 5% warmup   (epochs 0-249)
    #   - 70% stable  (epochs 250-3749)  ← loss can actually move here
    #   - 25% decay   (epochs 3750-4999)
    # CosineAnnealingLR decayed to near-zero by epoch 75, which is why training stalled.
    TOTAL_EPOCHS  = 5000
    WARMUP_EPOCHS = int(TOTAL_EPOCHS * 0.05)
    DECAY_START   = int(TOTAL_EPOCHS * 0.75)

    def wsd_lr(epoch):
        if epoch < WARMUP_EPOCHS:
            return epoch / max(WARMUP_EPOCHS, 1)
        if epoch < DECAY_START:
            return 1.0
        progress = (epoch - DECAY_START) / max(TOTAL_EPOCHS - DECAY_START, 1)
        # inverse-square-root decay as used in the H-Net paper
        return max(1.0 / math.sqrt(1.0 + 10 * progress), 0.05)

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, wsd_lr)

    # Resume from the latest checkpoint if one exists
    start_epoch = 0
    if RESUME_LATEST:
        ckpts = sorted([
            f for f in os.listdir(SAVE_DIR) if f.startswith("ckpt_epoch") and f.endswith(".pt")
        ])
        if ckpts:
            start_epoch = load_checkpoint(
                os.path.join(SAVE_DIR, ckpts[-1]), model, optimizer, scheduler
            )

    # ── Data loading ─────────────────────────────────────────────────────────────
    # Reads one or more parquet files and concatenates the text column into a
    # single flat byte tensor.  Expects HuggingFace Finnish Wikipedia format
    # (columns: id, url, title, text) — change TEXT_COL if yours differs.
    TEXT_COL     = "text"
    PARQUET_FILES = [
        "data/train-00000-of-00002.parquet",
        "data/train-00001-of-00002.parquet",
    ]

    print("Loading parquet data...", flush=True)
    try:
        import pyarrow.parquet as pq
        chunks = []
        for path in PARQUET_FILES:
            table = pq.read_table(path, columns=[TEXT_COL])
            for text in table.column(TEXT_COL).to_pylist():
                if text:
                    chunks.append(text.replace("\n", " "))
        full_text = " ".join(chunks)
        raw = torch.frombuffer(
            bytearray(full_text.encode("utf-8")), dtype=torch.uint8
        ).long()
        print(f"Training data   : {raw.shape[0]:,} bytes  ({len(chunks):,} articles)")
    except (ImportError, FileNotFoundError) as e:
        print(f"Parquet load failed ({e}), falling back to kalevala.plain.txt")
        with open("data/kalevala.plain.txt", "r", encoding="utf-8") as fd:
            text = fd.read().replace("\n", " ")
        raw = torch.frombuffer(
            bytearray(text.encode("utf-8")), dtype=torch.uint8
        ).long()
        print(f"Training data   : {raw.shape[0]:,} bytes")

    N_tok = raw.shape[0]

    def sample_batch() -> torch.Tensor:
        starts = torch.randint(0, N_tok - SEQ_LEN, (BATCH,))
        return torch.stack([raw[s : s + SEQ_LEN] for s in starts]).to(device)

    for epoch in range(start_epoch, TOTAL_EPOCHS):
        model.train()
        total_loss = 0.0

        # Curriculum: linearly ramp masked loss from 0% → 50% over 100 epochs.
        # Mask ratio itself also grows 0.3 → 0.5 over 200 epochs — start with
        # easy gaps (few masked positions) and get harder as the model improves.
        epochs_since_start = epoch - start_epoch
        masked_ratio = min(0.5, epochs_since_start / 100.0 * 0.5)
        mask_prob    = min(0.5, 0.3 + epochs_since_start / 200.0 * 0.2)

        for _ in range(200):
            optimizer.zero_grad(set_to_none=True)
            torch.cuda.empty_cache()
            batch = sample_batch()
            if torch.rand(1).item() < masked_ratio:
                # backward() called inside masked_diffusion_loss per sample
                loss = masked_diffusion_loss(
                    model, diffusion, batch,
                    mask_ratio=mask_prob,
                    rounding_weight=0.5,
                )
            else:
                loss = diffusion_loss(
                    model, diffusion, batch,
                    rounding_weight=0.5,
                    rounding_every=1,
                )
                loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_loss += loss.item()
            del loss, batch
        scheduler.step()

        avg_loss   = total_loss / 200
        current_lr = scheduler.get_last_lr()[0]
        print(f"epoch {epoch:4d}  loss {avg_loss:.4f}  lr {current_lr:.2e}  masked {masked_ratio:.0%}  mask_p {mask_prob:.2f}")

        model.eval()
        embs     = ar_sample(model, diffusion, n_steps=80, device=device)
        ids      = decode(model, embs)
        text_out = ids[0].byte().cpu().numpy().tobytes().decode("utf-8", errors="replace")
        print(f"  → {text_out!r}")

        if (epoch + 1) % SAVE_EVERY == 0:
            ckpt_path = os.path.join(SAVE_DIR, f"ckpt_epoch{epoch:04d}.pt")
            save_checkpoint(ckpt_path, epoch, model, optimizer, scheduler, avg_loss)