Blame for modeling_cola2.py

"""Hugging Face compatible Cola2 span-causal Full-DiT model."""

from __future__ import annotations

import copy
import json
import math
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn
from transformers import PreTrainedModel
from transformers.generation.utils import GenerateDecoderOnlyOutput
from transformers.modeling_outputs import ModelOutput

from .configuration_cola2 import Cola2Config, Cola2GenerationConfig
from .modeling_layers import (
    AdaLN,
    DiTBlock,
    DurationHead,
    HardDurationExpander,
    Qwen3DecoderBlock,
    RMSNorm,
    TextVAEBlock,
    TimestepEmbedding,
    additive_attention_mask,
    block_causal_allow,
    causal_padding_allow,
    span_local_allow,
)


def _build_dit_timesteps(
    *,
    schedule_t: float,
    steps: int,
    schedule: str,
    power: Optional[float],
    device: torch.device,
) -> torch.Tensor:
    named_powers = {
        "uniform": 1.0,
        "high_noise": 0.5,
        "low_noise": 2.0,
    }
    if schedule == "power":
        exponent = float(power if power is not None else 1.0)
    else:
        exponent = named_powers[schedule]
    progress = torch.linspace(
        1.0,
        0.0,
        int(steps) + 1,
        dtype=torch.float32,
        device=device,
    )
    return progress.pow(exponent).mul_(float(schedule_t))


@dataclass
class Cola2EncoderOutput(ModelOutput):
    """Dynamic compression encoder output."""

    latents: Optional[torch.Tensor] = None
    latent_logvar: Optional[torch.Tensor] = None
    latent_mask: Optional[torch.Tensor] = None
    span_starts: Optional[torch.Tensor] = None
    span_ends: Optional[torch.Tensor] = None
    span_lengths: Optional[torch.Tensor] = None
    segment_ids: Optional[torch.Tensor] = None
    boundary_logits: Optional[torch.Tensor] = None
    boundary_probabilities: Optional[torch.Tensor] = None


@dataclass
class Cola2ConditionalGenerationOutput(ModelOutput):
    """Output of the same-position span-causal reconstruction forward."""

    loss: Optional[torch.Tensor] = None
    logits: Optional[torch.Tensor] = None
    latents: Optional[torch.Tensor] = None
    latent_mask: Optional[torch.Tensor] = None
    span_lengths: Optional[torch.Tensor] = None
    duration_logits: Optional[torch.Tensor] = None
    boundary_logits: Optional[torch.Tensor] = None


class CosineBoundaryScorer(nn.Module):
    """Native cosine scorer layout (``segmenter.scorer.norm`` in checkpoints)."""

    def __init__(self, config: Cola2Config) -> None:
        super().__init__()
        self.alpha = config.boundary_score_alpha
        self.threshold = config.boundary_score_threshold
        self.norm = nn.LayerNorm(
            config.encoder_hidden_size, eps=config.layer_norm_eps
        )

    def forward(
        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        batch, width, _ = hidden_states.shape
        if width <= 1:
            return hidden_states.new_zeros((batch, 0))
        normalized = self.norm(hidden_states)
        cosine = F.cosine_similarity(
            normalized[:, :-1], normalized[:, 1:], dim=-1
        )
        logits = self.alpha * (1.0 - cosine - self.threshold)
        valid = attention_mask.bool()
        return logits.masked_fill(
            ~(valid[:, :-1] & valid[:, 1:]), -1e4
        )


class DynamicSegmenter(nn.Module):
    """Auxiliary-loss-free cosine boundary scorer with hard span constraints."""

    def __init__(self, config: Cola2Config) -> None:
        super().__init__()
        self.lmin = config.min_span_length
        self.lmax = config.max_span_length
        self.scorer = CosineBoundaryScorer(config)
        rate = min(max(config.target_compression_rate, 1e-4), 1.0 - 1e-4)
        self.register_buffer(
            "rate_bias",
            torch.tensor(math.log(rate / (1.0 - rate)), dtype=torch.float32),
        )

    def forward(
        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
    ) -> dict[str, torch.Tensor]:
        batch, width, _ = hidden_states.shape
        valid = attention_mask.bool()
        logits = self.scorer(hidden_states, valid)
        if width > 1:
            probabilities = torch.sigmoid(
                logits + self.rate_bias.to(logits.dtype)
            )
        else:
            logits = hidden_states.new_zeros((batch, 0))
            probabilities = hidden_states.new_zeros((batch, 0))

        rows: list[list[tuple[int, int]]] = []
        boundaries = torch.zeros(
            batch, width, dtype=torch.bool, device=hidden_states.device
        )
        for row in range(batch):
            length = int(valid[row].long().sum().item())
            if length <= 0:
                rows.append([])
                continue
            spans: list[tuple[int, int]] = []
            start = 0
            for token_index in range(length):
                span_length = token_index - start + 1
                wanted = (
                    token_index < length - 1
                    and bool(probabilities[row, token_index] > 0.5)
                )
                should_cut = (
                    token_index == length - 1
                    or span_length >= self.lmax
                    or (wanted and span_length >= self.lmin)
                )
                if should_cut:
                    spans.append((start, token_index))
                    boundaries[row, token_index] = True
                    start = token_index + 1
            rows.append(spans)

        max_spans = max((len(row) for row in rows), default=0)
        max_spans = max(max_spans, 1)
        starts = torch.zeros(
            batch, max_spans, dtype=torch.long, device=hidden_states.device
        )
        ends = torch.zeros_like(starts)
        lengths = torch.zeros_like(starts)
        latent_mask = torch.zeros(
            batch, max_spans, dtype=torch.bool, device=hidden_states.device
        )
        segment_ids = torch.full(
            (batch, width), -1, dtype=torch.long, device=hidden_states.device
        )
        for row, spans in enumerate(rows):
            for span_index, (start, end) in enumerate(spans):
                starts[row, span_index] = start
                ends[row, span_index] = end
                lengths[row, span_index] = end - start + 1
                latent_mask[row, span_index] = True
                segment_ids[row, start : end + 1] = span_index

        return {
            "boundary_logits": logits,
            "boundary_probabilities": probabilities,
            "boundary_hard": boundaries,
            "span_starts": starts,
            "span_ends": ends,
            "span_lengths": lengths,
            "latent_mask": latent_mask,
            "segment_ids": segment_ids,
        }


class SpanLocalPool(nn.Module):
    """Span-local bidirectional contextualization followed by mean pooling."""

    def __init__(self, config: Cola2Config) -> None:
        super().__init__()
        self.blocks = nn.ModuleList(
            TextVAEBlock(
                config.encoder_hidden_size,
                config.encoder_intermediate_size,
                config.encoder_num_attention_heads,
                norm_eps=config.layer_norm_eps,
                rope_theta=config.encoder_rope_theta,
                attention_dropout=config.attention_dropout,
                hidden_dropout=config.hidden_dropout,
                norm_style="post_legacy",
            )
            for _ in range(config.pool_num_hidden_layers)
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        segment_ids: torch.Tensor,
        latent_mask: torch.Tensor,
    ) -> torch.Tensor:
        valid = attention_mask.bool()
        positions = (torch.cumsum(valid.long(), dim=1) - 1).clamp_min(0)
        allowed = span_local_allow(valid, segment_ids)
        bias = additive_attention_mask(allowed, hidden_states.dtype)
        for block in self.blocks:
            hidden_states = block(hidden_states, bias, positions)

        batch, _, dim = hidden_states.shape
        max_spans = latent_mask.shape[1]
        pooled = hidden_states.new_zeros((batch, max_spans, dim))
        for row in range(batch):
            for span_index in range(max_spans):
                if not bool(latent_mask[row, span_index]):
                    continue
                selected = segment_ids[row] == span_index
                pooled[row, span_index] = hidden_states[row, selected].mean(
                    dim=0
                )
        return pooled


class Cola2CompressionVAE(nn.Module):
    """Dynamic text compressor and span-causal token decoder."""

    def __init__(self, config: Cola2Config) -> None:
        super().__init__()
        self.config = config
        input_rows = max(
            config.vocab_size,
            int(config.pad_token_id or 0) + 1,
            int(config.mask_token_id or 0) + 1,
        )
        self.encoder = nn.ModuleDict(
            {
                "wte": nn.Embedding(input_rows, config.encoder_embedding_size),
                "input_proj": (
                    nn.Linear(
                        config.encoder_embedding_size,
                        config.encoder_hidden_size,
                        bias=False,
                    )
                    if config.encoder_embedding_size
                    != config.encoder_hidden_size
                    else nn.Identity()
                ),
                "blocks": nn.ModuleList(
                    TextVAEBlock(
                        config.encoder_hidden_size,
                        config.encoder_intermediate_size,
                        config.encoder_num_attention_heads,
                        norm_eps=config.layer_norm_eps,
                        rope_theta=config.encoder_rope_theta,
                        attention_dropout=config.attention_dropout,
                        hidden_dropout=config.hidden_dropout,
                        norm_style="post_legacy",
                    )
                    for _ in range(config.encoder_num_hidden_layers)
                ),
            }
        )
        self.segmenter = DynamicSegmenter(config)
        self.pool = SpanLocalPool(config)
        self.latent_encoder_blocks = nn.ModuleList(
            TextVAEBlock(
                config.encoder_hidden_size,
                config.encoder_intermediate_size,
                config.encoder_num_attention_heads,
                norm_eps=config.layer_norm_eps,
                rope_theta=config.encoder_rope_theta,
                attention_dropout=config.attention_dropout,
                hidden_dropout=config.hidden_dropout,
                norm_style="post_legacy",
            )
            for _ in range(config.latent_encoder_num_hidden_layers)
        )
        posterior_size = config.latent_dim * (
            2 if config.use_variational_encoder else 1
        )
        self.posterior_proj = nn.Linear(
            config.encoder_hidden_size, posterior_size, bias=True
        )
        self.posterior_norm = nn.LayerNorm(
            config.latent_dim,
            eps=config.layer_norm_eps,
            elementwise_affine=False,
        )

        self.latent_in_proj = nn.Linear(
            config.latent_dim, config.decoder_hidden_size, bias=True
        )
        condition_intermediate = config.decoder_intermediate_size * 2
        condition_heads = config.condition_num_attention_heads
        self.latent_decoder_blocks = nn.ModuleList(
            TextVAEBlock(
                config.decoder_hidden_size,
                condition_intermediate,
                condition_heads,
                norm_eps=config.layer_norm_eps,
                rope_theta=config.condition_rope_theta,
                attention_dropout=config.attention_dropout,
                hidden_dropout=config.hidden_dropout,
                norm_style="reordered",
            )
            for _ in range(config.latent_decoder_num_hidden_layers)
        )
        self.duration_head = DurationHead(
            config.decoder_hidden_size,
            config.duration_hidden_size,
            config.max_span_length,
        )
        self.expander = HardDurationExpander(
            config.decoder_hidden_size,
            config.decoder_hidden_size,
            config.local_position_dim,
        )
        self.decoder_cond_stack = nn.ModuleList(
            TextVAEBlock(
                config.decoder_hidden_size,
                condition_intermediate,
                condition_heads,
                norm_eps=config.layer_norm_eps,
                rope_theta=config.condition_rope_theta,
                attention_dropout=config.attention_dropout,
                hidden_dropout=config.hidden_dropout,
                norm_style="reordered",
            )
            for _ in range(config.decoder_condition_num_hidden_layers)
        )

        self.decoder = nn.ModuleDict(
            {
                "blocks": nn.ModuleList(
                    Qwen3DecoderBlock(
                        config.decoder_hidden_size,
                        config.decoder_intermediate_size,
                        config.decoder_num_attention_heads,
                        config.decoder_num_key_value_heads,
                        config.decoder_head_dim,
                        norm_eps=config.layer_norm_eps,
                        rope_theta=config.decoder_rope_theta,
                        attention_dropout=config.attention_dropout,
                        hidden_dropout=config.hidden_dropout,
                    )
                    for _ in range(config.decoder_num_hidden_layers)
                ),
                "final_norm": RMSNorm(
                    config.decoder_hidden_size, eps=config.layer_norm_eps
                ),
                "final_layer": nn.Linear(
                    config.decoder_hidden_size,
                    config.vocab_size,
                    bias=config.decoder_lm_head_bias,
                ),
            }
        )
        self.decoder_mask_embedding = nn.Parameter(
            torch.zeros(config.decoder_hidden_size)
        )
        self.decoder_cond_in = nn.Linear(
            config.decoder_hidden_size,
            config.decoder_hidden_size,
            bias=True,
        )

    def token_hidden(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        hidden_states = self.encoder["input_proj"](
            self.encoder["wte"](input_ids)
        )
        valid = attention_mask.bool()
        positions = (torch.cumsum(valid.long(), dim=1) - 1).clamp_min(0)
        bias = additive_attention_mask(
            causal_padding_allow(valid), hidden_states.dtype
        )
        for block in self.encoder["blocks"]:
            hidden_states = block(hidden_states, bias, positions)
        return hidden_states

    def _latent_encode(
        self, pooled: torch.Tensor, latent_mask: torch.Tensor
    ) -> torch.Tensor:
        if not self.latent_encoder_blocks:
            return pooled
        valid = latent_mask.bool()
        positions = torch.arange(
            pooled.shape[1], device=pooled.device, dtype=torch.long
        ).unsqueeze(0).expand(pooled.shape[0], -1)
        if self.config.latent_block_causal:
            block_ids = positions // self.config.latent_block_size
            allowed = block_causal_allow(valid, block_ids)
        else:
            allowed = causal_padding_allow(valid)
        bias = additive_attention_mask(allowed, pooled.dtype)
        for block in self.latent_encoder_blocks:
            pooled = block(pooled, bias, positions)
        return pooled

    def _posterior(
        self, hidden_states: torch.Tensor
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        parameters = self.posterior_proj(hidden_states)
        if self.config.use_variational_encoder:
            mean, logvar = parameters.chunk(2, dim=-1)
            return self.posterior_norm(mean), logvar.clamp(-30.0, 20.0)
        return self.posterior_norm(parameters), None

    def encode_dynamic(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> Cola2EncoderOutput:
        hidden_states = self.token_hidden(input_ids, attention_mask)
        segmentation = self.segmenter(hidden_states, attention_mask)
        pooled = self.pool(
            hidden_states,
            attention_mask,
            segmentation["segment_ids"],
            segmentation["latent_mask"],
        )
        pooled = self._latent_encode(pooled, segmentation["latent_mask"])
        latents, logvar = self._posterior(pooled)
        latents = latents * segmentation["latent_mask"].unsqueeze(-1)
        return Cola2EncoderOutput(
            latents=latents,
            latent_logvar=logvar,
            latent_mask=segmentation["latent_mask"],
            span_starts=segmentation["span_starts"],
            span_ends=segmentation["span_ends"],
            span_lengths=segmentation["span_lengths"],
            segment_ids=segmentation["segment_ids"],
            boundary_logits=segmentation["boundary_logits"],
            boundary_probabilities=segmentation["boundary_probabilities"],
        )

    def encode_given_spans(
        self, input_ids: torch.Tensor, span_lengths: torch.Tensor
    ) -> torch.Tensor:
        """Re-encode committed tokens under decoder-predicted span geometry."""

        if input_ids.ndim != 1 or span_lengths.ndim != 1:
            raise ValueError("encode_given_spans expects unpadded 1-D rows.")
        if int(span_lengths.sum().item()) != int(input_ids.numel()):
            raise ValueError(
                "Feedback spans must be positive and cover every committed token."
            )
        if bool((span_lengths <= 0).any()):
            raise ValueError("Feedback span lengths must all be positive.")
        token_batch = input_ids.unsqueeze(0)
        valid = torch.ones_like(token_batch, dtype=torch.bool)
        hidden_states = self.token_hidden(token_batch, valid)
        segment_ids = torch.empty_like(token_batch)
        offset = 0
        for span_index, length_tensor in enumerate(span_lengths):
            length = int(length_tensor.item())
            segment_ids[:, offset : offset + length] = span_index
            offset += length
        latent_mask = torch.ones(
            1,
            span_lengths.numel(),
            dtype=torch.bool,
            device=input_ids.device,
        )
        pooled = self.pool(
            hidden_states, valid, segment_ids, latent_mask
        )
        pooled = self._latent_encode(pooled, latent_mask)
        latents, _ = self._posterior(pooled)
        return latents[0]

    def latent_decode(
        self, latents: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if latents.ndim != 2:
            raise ValueError("latent_decode expects [K, latent_dim].")
        hidden_states = self.latent_in_proj(latents).unsqueeze(0)
        length = hidden_states.shape[1]
        valid = torch.ones(
            1, length, dtype=torch.bool, device=latents.device
        )
        positions = torch.arange(
            length, device=latents.device, dtype=torch.long
        ).unsqueeze(0)
        if self.config.latent_block_causal:
            block_ids = positions // self.config.latent_block_size
            allowed = block_causal_allow(valid, block_ids)
        else:
            allowed = causal_padding_allow(valid)
        bias = additive_attention_mask(allowed, hidden_states.dtype)
        for block in self.latent_decoder_blocks:
            hidden_states = block(hidden_states, bias, positions)
        duration_logits = self.duration_head(hidden_states)
        return hidden_states[0], duration_logits[0]

    def expand_and_refine(
        self, latents: torch.Tensor, span_lengths: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Create token-aligned latent conditions for one complete sequence."""

        latent_hidden, duration_logits = self.latent_decode(latents)
        condition = self.expander(latent_hidden, span_lengths)
        if condition.shape[0] == 0:
            return condition, duration_logits
        latent_blocks = (
            torch.arange(
                span_lengths.numel(),
                device=latents.device,
                dtype=torch.long,
            )
            // self.config.dit_block_size
        )
        token_blocks = torch.repeat_interleave(
            latent_blocks, span_lengths.long(), dim=0
        ).unsqueeze(0)
        valid = torch.ones(
            1, condition.shape[0], dtype=torch.bool, device=latents.device
        )
        positions = torch.arange(
            condition.shape[0], device=latents.device, dtype=torch.long
        ).unsqueeze(0)
        bias = additive_attention_mask(
            block_causal_allow(valid, token_blocks), condition.dtype
        )
        refined = condition.unsqueeze(0)
        for block in self.decoder_cond_stack:
            refined = block(refined, bias, positions)
        return refined[0], duration_logits

    def span_query_logits(
        self,
        clean_ids: torch.Tensor,
        query_condition: torch.Tensor,
        *,
        conditional: bool,
    ) -> torch.Tensor:
        """Predict one MASK-query chunk against the clean token prefix."""

        if clean_ids.ndim != 1 or query_condition.ndim != 2:
            raise ValueError("span_query_logits expects one unpadded row.")
        query_width = query_condition.shape[0]
        if query_width == 0:
            return query_condition.new_zeros((0, self.config.vocab_size))
        output_weight = self.decoder["final_layer"].weight
        if clean_ids.numel():
            clean_embedding = F.embedding(
                clean_ids.clamp(0, self.config.vocab_size - 1),
                output_weight,
            )
        else:
            clean_embedding = query_condition.new_zeros(
                (0, self.config.decoder_hidden_size)
            )
        query_embedding = self.decoder_mask_embedding.unsqueeze(0).expand(
            query_width, -1
        )
        if conditional:
            query_embedding = query_embedding + self.decoder_cond_in(
                query_condition
            )
        hidden_states = torch.cat(
            (clean_embedding, query_embedding), dim=0
        ).unsqueeze(0)
        total = hidden_states.shape[1]
        valid = torch.ones(
            1, total, dtype=torch.bool, device=hidden_states.device
        )
        positions = torch.arange(
            total, device=hidden_states.device, dtype=torch.long
        ).unsqueeze(0)
        bias = additive_attention_mask(
            causal_padding_allow(valid), hidden_states.dtype
        )
        for block in self.decoder["blocks"]:
            hidden_states = block(hidden_states, bias, positions)
        query_states = hidden_states[:, -query_width:]
        query_states = self.decoder["final_norm"](query_states)
        return self.decoder["final_layer"](query_states)[0]

    def teacher_forced_logits(
        self,
        input_ids: torch.Tensor,
        latents: torch.Tensor,
        span_lengths: torch.Tensor,
        *,
        chunk_spans: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Same-position dual-stream logits used by :meth:`forward`."""

        condition, duration_logits = self.expand_and_refine(
            latents, span_lengths
        )
        logits = condition.new_zeros(
            input_ids.numel(), self.config.vocab_size
        )
        bounds = torch.cat(
            (
                span_lengths.new_zeros(1),
                torch.cumsum(span_lengths, dim=0),
            ),
            dim=0,
        )
        for span_lo in range(0, span_lengths.numel(), chunk_spans):
            span_hi = min(span_lo + chunk_spans, span_lengths.numel())
            token_lo = int(bounds[span_lo].item())
            token_hi = int(bounds[span_hi].item())
            logits[token_lo:token_hi] = self.span_query_logits(
                input_ids[:token_lo],
                condition[token_lo:token_hi],
                conditional=True,
            )
        return logits, duration_logits


class _PatchProjection(nn.Module):
    def __init__(self, input_size: int, output_size: int) -> None:
        super().__init__()
        self.proj = nn.Linear(input_size, output_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.proj(hidden_states)


class Cola2DiT(nn.Module):
    """Block-causal latent flow transformer."""

    def __init__(self, config: Cola2Config) -> None:
        super().__init__()
        self.block_size = config.dit_block_size
        self.txt_in = _PatchProjection(
            config.latent_dim, config.dit_hidden_size
        )
        self.emb_in = TimestepEmbedding(
            config.dit_hidden_size, config.dit_hidden_size
        )
        self.blocks = nn.ModuleList(
            DiTBlock(
                config.dit_hidden_size,
                config.dit_intermediate_size,
                config.dit_num_attention_heads,
                config.dit_head_dim,
                config.dit_rope_dim,
                norm_eps=config.dit_layer_norm_eps,
                rope_theta=config.dit_rope_theta,
                attention_dropout=config.attention_dropout,
            )
            for _ in range(config.dit_num_hidden_layers)
        )
        self.txt_out_norm = nn.LayerNorm(
            config.dit_hidden_size,
            eps=config.dit_layer_norm_eps,
            elementwise_affine=True,
        )
        self.txt_out_ada = AdaLN(config.dit_hidden_size, ("out",))
        self.txt_out = _PatchProjection(
            config.dit_hidden_size, config.latent_dim
        )

    def forward(
        self, latent_states: torch.Tensor, timesteps: torch.Tensor
    ) -> torch.Tensor:
        if latent_states.ndim != 2:
            raise ValueError("DiT expects one latent row [L,D].")
        length = latent_states.shape[0]
        if timesteps.ndim == 0:
            timesteps = timesteps.expand(length)
        if timesteps.numel() == 1:
            timesteps = timesteps.expand(length)
        if timesteps.numel() != length:
            raise ValueError("DiT needs one timestep per latent position.")
        hidden_states = self.txt_in(latent_states).unsqueeze(0)
        embedding = self.emb_in(timesteps).unsqueeze(0)
        positions = torch.arange(
            length, device=latent_states.device, dtype=torch.long
        ).unsqueeze(0)
        valid = torch.ones(
            1, length, dtype=torch.bool, device=latent_states.device
        )
        block_ids = positions // self.block_size
        bias = additive_attention_mask(
            block_causal_allow(valid, block_ids), hidden_states.dtype
        )
        for block in self.blocks:
            hidden_states = block(
                hidden_states, embedding, bias, positions
            )
        hidden_states = self.txt_out_ada.modulate(
            hidden_states,
            embedding,
            "out",
            self.txt_out_norm,
        )
        return self.txt_out(hidden_states)[0]


class Cola2ForConditionalGeneration(PreTrainedModel):
    """Runnable span-causal Cola2 Full-DiT generation wrapper.

    Generation follows the production sequence:

    ``prompt encode -> DiT denoise -> dynamic duration -> span-causal decode
    -> feedback re-encode``.
    """

    config_class = Cola2Config
    generation_config_class = Cola2GenerationConfig
    base_model_prefix = "cola2"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = False
    _no_split_modules = ["Qwen3DecoderBlock", "DiTBlock", "TextVAEBlock"]

    def __init__(self, config: Cola2Config) -> None:
        super().__init__(config)
        self.vae = Cola2CompressionVAE(config)
        self.dit = Cola2DiT(config)
        self.last_generation_stats: dict[str, int] = {}
        self.post_init()
        self.generation_config = Cola2GenerationConfig(
            pad_token_id=config.pad_token_id,
            bos_token_id=config.bos_token_id,
            eos_token_id=config.eos_token_id,
        )

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(
                module.weight, mean=0.0, std=self.config.initializer_range
            )
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(
                module.weight, mean=0.0, std=self.config.initializer_range
            )
        elif isinstance(module, nn.LayerNorm):
            if module.weight is not None:
                nn.init.ones_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, RMSNorm):
            nn.init.ones_(module.weight)

    @classmethod
    def from_pretrained(cls, *args: Any, **kwargs: Any) -> Any:
        """Preserve Cola2-specific generation fields after HF loading."""

        result = super().from_pretrained(*args, **kwargs)
        model = result[0] if isinstance(result, tuple) else result
        if not isinstance(model.generation_config, Cola2GenerationConfig):
            model.generation_config = Cola2GenerationConfig.from_dict(
                model.generation_config.to_dict()
            )
        return result

    def get_input_embeddings(self) -> nn.Module:
        return self.vae.encoder["wte"]

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.vae.encoder["wte"] = value

    def get_output_embeddings(self) -> nn.Module:
        return self.vae.decoder["final_layer"]

    def set_output_embeddings(self, value: nn.Module) -> None:
        self.vae.decoder["final_layer"] = value

    def encode(
        self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
    ) -> Cola2EncoderOutput:
        if attention_mask is None:
            attention_mask = input_ids.ne(self.config.pad_token_id)
        return self.vae.encode_dynamic(input_ids, attention_mask)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
        **_: Any,
    ) -> Union[Cola2ConditionalGenerationOutput, tuple[torch.Tensor, ...]]:
        """Encode and run the real span-causal reconstruction graph.

        Unlike an autoregressive LM, Cola2 predicts each token at the same
        position from its latent condition and the clean prefix available at the
        start of that complete-span chunk.
        """

        return_dict = (
            self.config.use_return_dict
            if return_dict is None
            else bool(return_dict)
        )
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence].")
        if attention_mask is None:
            attention_mask = input_ids.ne(self.config.pad_token_id)
        attention_mask = attention_mask.bool()
        encoder_output = self.vae.encode_dynamic(input_ids, attention_mask)
        batch, width = input_ids.shape
        logits = next(self.parameters()).new_zeros(
            batch, width, self.config.vocab_size
        )
        duration_logits = next(self.parameters()).new_zeros(
            batch,
            encoder_output.latents.shape[1],
            self.config.max_span_length,
        )
        chunk_spans = int(
            getattr(self.generation_config, "span_causal_chunk_size", 1)
        )
        for row in range(batch):
            token_count = int(attention_mask[row].long().sum().item())
            latent_count = int(
                encoder_output.latent_mask[row].long().sum().item()
            )
            if token_count <= 0 or latent_count <= 0:
                continue
            row_logits, row_duration = self.vae.teacher_forced_logits(
                input_ids[row, attention_mask[row]],
                encoder_output.latents[row, :latent_count],
                encoder_output.span_lengths[row, :latent_count],
                chunk_spans=chunk_spans,
            )
            logits[row, attention_mask[row]] = row_logits
            duration_logits[row, :latent_count] = row_duration

        loss = None
        if labels is not None:
            if labels.shape != input_ids.shape:
                raise ValueError("labels must have the same shape as input_ids.")
            supervised = attention_mask & labels.ge(0)
            if bool(supervised.any()):
                loss = F.cross_entropy(
                    logits[supervised].float(),
                    labels[supervised].clamp(
                        min=0, max=self.config.vocab_size - 1
                    ),
                )
            else:
                loss = logits.sum() * 0.0

        output = Cola2ConditionalGenerationOutput(
            loss=loss,
            logits=logits,
            latents=encoder_output.latents,
            latent_mask=encoder_output.latent_mask,
            span_lengths=encoder_output.span_lengths,
            duration_logits=duration_logits,
            boundary_logits=encoder_output.boundary_logits,
        )
        if return_dict:
            return output
        values: tuple[torch.Tensor, ...] = (
            logits,
            encoder_output.latents,
            encoder_output.span_lengths,
            duration_logits,
        )
        return ((loss,) + values) if loss is not None else values

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            **kwargs,
        }

    def _scale_latents(self, latents: torch.Tensor) -> torch.Tensor:
        return (
            latents - self.config.shifting_factor
        ) * self.config.scaling_factor

    def _unscale_latents(self, latents: torch.Tensor) -> torch.Tensor:
        return (
            latents / self.config.scaling_factor
        ) + self.config.shifting_factor

    def _dit_drift(
        self,
        prefix: torch.Tensor,
        block: torch.Tensor,
        block_timesteps: torch.Tensor,
        *,
        conditional: bool,
    ) -> torch.Tensor:
        if conditional and prefix.numel():
            states = torch.cat((prefix, block), dim=0)
            timesteps = torch.cat(
                (
                    block_timesteps.new_zeros(prefix.shape[0]),
                    block_timesteps,
                ),
                dim=0,
            )
            return self.dit(states, timesteps)[-block.shape[0] :]
        return self.dit(block, block_timesteps)

    def _denoise_block(
        self,
        prefix: torch.Tensor,
        block_init: torch.Tensor,
        clean_prefix_length: int,
        generation_config: Cola2GenerationConfig,
    ) -> torch.Tensor:
        state = block_init.float()
        clean = block_init[:clean_prefix_length].float().clone()
        timesteps = _build_dit_timesteps(
            schedule_t=self.config.diffusion_time,
            steps=generation_config.dit_num_steps,
            schedule=generation_config.dit_timestep_schedule,
            power=generation_config.dit_timestep_power,
            device=state.device,
        )
        model_dtype = self.dit.txt_in.proj.weight.dtype
        for current, following in zip(
            timesteps[:-1], timesteps[1:], strict=True
        ):
            step = torch.full(
                (state.shape[0],),
                float(current.item()),
                dtype=torch.float32,
                device=state.device,
            )
            if clean_prefix_length:
                step[:clean_prefix_length] = 0.0
            state_model = state.to(model_dtype)
            prefix_model = prefix.to(model_dtype)
            step_model = step.to(model_dtype)
            conditional = self._dit_drift(
                prefix_model,
                state_model,
                step_model,
                conditional=True,
            ).float()
            apply_cfg = (
                abs(generation_config.dit_guidance_scale - 1.0) > 1e-8
                and not (
                    generation_config.dit_cfg_disable_empty_prefix
                    and prefix.numel() == 0
                )
            )
            if apply_cfg:
                unconditional = self._dit_drift(
                    prefix_model,
                    state_model,
                    step_model,
                    conditional=False,
                ).float()
                drift = (
                    generation_config.dit_guidance_scale
                    * (conditional - unconditional)
                    + unconditional
                )
            else:
                drift = conditional
            dt = float((current - following).item()) / self.config.diffusion_time
            state = state - drift * dt
            if clean_prefix_length:
                state[:clean_prefix_length] = clean
            self.last_generation_stats["dit_denoise_steps"] += 1
        return state

    @staticmethod
    def _filter_logits(
        logits: torch.Tensor, top_k: int, top_p: float
    ) -> torch.Tensor:
        filtered = logits.clone()
        if top_k > 0:
            k = min(int(top_k), filtered.shape[-1])
            threshold = torch.topk(filtered, k, dim=-1).values[:, -1:]
            filtered.masked_fill_(filtered < threshold, float("-inf"))
        if 0.0 < top_p < 1.0:
            sorted_logits, sorted_indices = filtered.sort(
                dim=-1, descending=True
            )
            cumulative = sorted_logits.softmax(-1).cumsum(-1)
            remove = cumulative > top_p
            remove[:, 1:] = remove[:, :-1].clone()
            remove[:, 0] = False
            remove = torch.zeros_like(remove).scatter(
                1, sorted_indices, remove
            )
            filtered.masked_fill_(remove, float("-inf"))
        return filtered

    def _sample_chunk(
        self,
        logits: torch.Tensor,
        clean_ids: torch.Tensor,
        generation_config: Cola2GenerationConfig,
        generator: Optional[torch.Generator],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        scores = logits.float().clone()
        penalty = float(generation_config.repetition_penalty)
        if penalty != 1.0 and clean_ids.numel():
            token_ids = clean_ids.unique().clamp(
                0, self.config.vocab_size - 1
            )
            selected = scores[:, token_ids]
            selected = torch.where(
                selected < 0, selected * penalty, selected / penalty
            )
            scores[:, token_ids] = selected
        if not generation_config.do_sample:
            return scores.argmax(dim=-1), scores
        temperature = float(generation_config.temperature)
        if temperature <= 0:
            raise ValueError("temperature must be positive when do_sample=True.")
        scores = scores / temperature
        scores = self._filter_logits(
            scores,
            int(generation_config.top_k),
            float(generation_config.top_p),
        )
        probabilities = scores.softmax(dim=-1)
        sampled = torch.multinomial(
            probabilities, num_samples=1, generator=generator
        ).squeeze(-1)
        return sampled, scores

    def _span_causal_sample(
        self,
        clean_ids: torch.Tensor,
        query_condition: torch.Tensor,
        span_lengths: torch.Tensor,
        generation_config: Cola2GenerationConfig,
        generator: Optional[torch.Generator],
    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
        committed = clean_ids
        sampled_pieces: list[torch.Tensor] = []
        score_rows: list[torch.Tensor] = []
        bounds = torch.cat(
            (
                span_lengths.new_zeros(1),
                torch.cumsum(span_lengths, dim=0),
            ),
            dim=0,
        )
        chunk_spans = generation_config.span_causal_chunk_size
        for span_lo in range(0, span_lengths.numel(), chunk_spans):
            span_hi = min(span_lo + chunk_spans, span_lengths.numel())
            token_lo = int(bounds[span_lo].item())
            token_hi = int(bounds[span_hi].item())
            condition = query_condition[token_lo:token_hi]
            logits_conditional = self.vae.span_query_logits(
                committed, condition, conditional=True
            ).float()
            weight = generation_config.decoder_guidance_scale
            if abs(weight) > 1e-12:
                logits_unconditional = self.vae.span_query_logits(
                    committed, condition, conditional=False
                ).float()
                logits = (
                    (1.0 + weight) * logits_conditional
                    - weight * logits_unconditional
                )
            else:
                logits = logits_conditional
            sampled, scores = self._sample_chunk(
                logits, committed, generation_config, generator
            )
            committed = torch.cat((committed, sampled), dim=0)
            sampled_pieces.append(sampled)
            score_rows.extend(scores.unbind(dim=0))
            self.last_generation_stats["decoder_chunks"] += 1
        if not sampled_pieces:
            return clean_ids.new_zeros((0,)), score_rows
        return torch.cat(sampled_pieces, dim=0), score_rows

    def _generate_one(
        self,
        prompt_ids: torch.Tensor,
        max_new_tokens: int,
        generation_config: Cola2GenerationConfig,
        generator: Optional[torch.Generator],
    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
        if prompt_ids.numel() == 0:
            raise ValueError("Cola2 generation requires a non-empty prompt.")
        attention_mask = torch.ones_like(prompt_ids, dtype=torch.bool)
        encoded = self.vae.encode_dynamic(
            prompt_ids.unsqueeze(0), attention_mask.unsqueeze(0)
        )
        latent_count = int(encoded.latent_mask[0].long().sum().item())
        current_unscaled = encoded.latents[0, :latent_count].float()
        current_lengths = encoded.span_lengths[0, :latent_count].long()
        committed_ids = prompt_ids.clone()
        visible_tokens: list[int] = []
        visible_scores: list[torch.Tensor] = []
        eos_token_ids = generation_config.eos_token_id
        if eos_token_ids is None:
            eos_set: set[int] = set()
        elif isinstance(eos_token_ids, int):
            eos_set = {int(eos_token_ids)}
        else:
            eos_set = {int(value) for value in eos_token_ids}

        while len(visible_tokens) < max_new_tokens:
            scaled = self._scale_latents(current_unscaled)
            block_size = self.config.dit_block_size
            remainder = int(scaled.shape[0]) % block_size
            if remainder:
                prefix_length = int(scaled.shape[0]) - remainder
                clean_block = scaled[-remainder:]
                new_latent_count = block_size - remainder
            else:
                prefix_length = int(scaled.shape[0])
                clean_block = scaled.new_zeros((0, scaled.shape[-1]))
                new_latent_count = block_size
            prefix = scaled[:prefix_length]
            noise = torch.randn(
                new_latent_count,
                self.config.latent_dim,
                dtype=torch.float32,
                device=scaled.device,
                generator=generator,
            )
            block_init = torch.cat((clean_block.float(), noise), dim=0)
            denoised = self._denoise_block(
                prefix,
                block_init,
                remainder,
                generation_config,
            )
            block_unscaled = self._unscale_latents(denoised)
            candidate = torch.cat(
                (current_unscaled[:prefix_length], block_unscaled), dim=0
            )

            _, all_duration_logits = self.vae.latent_decode(
                candidate.to(self.vae.latent_in_proj.weight.dtype)
            )
            self.last_generation_stats["duration_predictions"] += 1
            if generation_config.dynamic_duration:
                predicted = all_duration_logits.argmax(dim=-1).long() + 1
            else:
                fixed = max(
                    1,
                    int(
                        round(
                            1.0 / self.config.target_compression_rate
                        )
                    ),
                )
                predicted = torch.full(
                    (candidate.shape[0],),
                    min(fixed, self.config.max_span_length),
                    dtype=torch.long,
                    device=candidate.device,
                )
            if remainder:
                predicted[prefix_length : prefix_length + remainder] = (
                    current_lengths[-remainder:]
                )
            new_lengths = predicted[-new_latent_count:]
            candidate_lengths = torch.cat(
                (current_lengths, new_lengths), dim=0
            )
            condition, _ = self.vae.expand_and_refine(
                candidate.to(self.vae.latent_in_proj.weight.dtype),
                candidate_lengths,
            )
            new_token_offset = int(current_lengths.sum().item())
            new_condition = condition[new_token_offset:]
            sampled, sampled_scores = self._span_causal_sample(
                committed_ids,
                new_condition,
                new_lengths,
                generation_config,
                generator,
            )
            committed_ids = torch.cat((committed_ids, sampled), dim=0)
            current_lengths = candidate_lengths

            stop = False
            for token, score in zip(
                sampled.detach().tolist(), sampled_scores, strict=True
            ):
                if len(visible_tokens) >= max_new_tokens:
                    break
                visible_tokens.append(int(token))
                visible_scores.append(score)
                if int(token) in eos_set:
                    stop = True
                    break

            if generation_config.latent_feedback:
                current_unscaled = self.vae.encode_given_spans(
                    committed_ids, current_lengths
                ).float()
                self.last_generation_stats["feedback_rounds"] += 1
            else:
                current_unscaled = candidate.float()
            self.last_generation_stats["generation_rounds"] += 1
            if stop:
                break

        return torch.tensor(
            visible_tokens, dtype=torch.long, device=prompt_ids.device
        ), visible_scores

    def _resolve_generation_config(
        self,
        generation_config: Optional[Union[Cola2GenerationConfig, Mapping[str, Any]]],
        kwargs: dict[str, Any],
    ) -> Cola2GenerationConfig:
        if generation_config is None:
            config = copy.deepcopy(self.generation_config)
        elif isinstance(generation_config, Mapping):
            config = Cola2GenerationConfig.from_dict(dict(generation_config))
        elif isinstance(generation_config, Cola2GenerationConfig):
            config = copy.deepcopy(generation_config)
        else:
            config = Cola2GenerationConfig.from_dict(
                generation_config.to_dict()
            )
        unused = config.update(**kwargs)
        if unused:
            names = ", ".join(sorted(unused))
            raise ValueError(f"Unused generation arguments: {names}.")
        config.validate()
        if int(config.num_return_sequences) != 1:
            raise ValueError("Cola2 currently supports num_return_sequences=1.")
        if int(config.num_beams) != 1:
            raise ValueError(
                "Beam search is not defined for block-parallel span-causal decoding."
            )
        return config

    @torch.no_grad()
    def generate(
        self,
        input_ids: Optional[torch.Tensor] = None,
        generation_config: Optional[
            Union[Cola2GenerationConfig, Mapping[str, Any]]
        ] = None,
        attention_mask: Optional[torch.Tensor] = None,
        generator: Optional[torch.Generator] = None,
        streamer: Any = None,
        **kwargs: Any,
    ) -> Union[torch.Tensor, GenerateDecoderOnlyOutput]:
        """Generate continuations with Full-DiT and span-causal decoding."""

        if input_ids is None:
            bos = self.config.bos_token_id
            if bos is None:
                raise ValueError("input_ids or bos_token_id is required.")
            input_ids = torch.tensor(
                [[bos]], dtype=torch.long, device=self.device
            )
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence].")
        if attention_mask is None:
            attention_mask = input_ids.ne(self.config.pad_token_id)
        if attention_mask.shape != input_ids.shape:
            raise ValueError("attention_mask must match input_ids.")

        config = self._resolve_generation_config(generation_config, kwargs)
        self.last_generation_stats = {
            "prompt_encodes": int(input_ids.shape[0]),
            "dit_denoise_steps": 0,
            "duration_predictions": 0,
            "decoder_chunks": 0,
            "feedback_rounds": 0,
            "generation_rounds": 0,
        }
        generated_rows: list[torch.Tensor] = []
        score_rows: list[list[torch.Tensor]] = []
        prompt_lengths: list[int] = []
        for row in range(input_ids.shape[0]):
            valid_ids = input_ids[row, attention_mask[row].bool()]
            prompt_lengths.append(int(valid_ids.numel()))
            if config.max_new_tokens is not None:
                budget = int(config.max_new_tokens)
            else:
                budget = max(0, int(config.max_length) - int(valid_ids.numel()))
            generated, scores = self._generate_one(
                valid_ids, budget, config, generator
            )
            generated_rows.append(generated)
            score_rows.append(scores)

        max_generated = max((row.numel() for row in generated_rows), default=0)
        pad_id = config.pad_token_id
        if pad_id is None:
            pad_id = (
                config.eos_token_id
                if isinstance(config.eos_token_id, int)
                else self.config.pad_token_id
            )
        if pad_id is None:
            pad_id = 0
        sequences = input_ids.new_full(
            (input_ids.shape[0], input_ids.shape[1] + max_generated),
            int(pad_id),
        )
        sequences[:, : input_ids.shape[1]] = input_ids
        for row, generated in enumerate(generated_rows):
            sequences[
                row,
                input_ids.shape[1] : input_ids.shape[1] + generated.numel(),
            ] = generated

        if streamer is not None:
            if input_ids.shape[0] != 1:
                raise ValueError("streamer is supported only for batch size 1.")
            for token in generated_rows[0]:
                streamer.put(token.reshape(1, 1).cpu())
            streamer.end()

        if not config.return_dict_in_generate:
            return sequences
        scores_tuple = None
        if config.output_scores:
            score_steps: list[torch.Tensor] = []
            for step in range(max_generated):
                rows = []
                for row_scores in score_rows:
                    if step < len(row_scores):
                        rows.append(row_scores[step])
                    else:
                        filler = next(self.parameters()).new_full(
                            (self.config.vocab_size,), float("-inf")
                        )
                        if 0 <= int(pad_id) < self.config.vocab_size:
                            filler[int(pad_id)] = 0.0
                        rows.append(filler)
                score_steps.append(torch.stack(rows, dim=0))
            scores_tuple = tuple(score_steps)
        return GenerateDecoderOnlyOutput(
            sequences=sequences,
            scores=scores_tuple,
        )

    @staticmethod
    def _read_checkpoint_state(
        source: Union[str, Path, Mapping[str, torch.Tensor]],
        map_location: Union[str, torch.device] = "cpu",
    ) -> dict[str, torch.Tensor]:
        if isinstance(source, Mapping):
            raw: Any = dict(source)
        else:
            path = Path(source)
            if path.is_dir():
                safe_index = path / "model.safetensors.index.json"
                bin_index = path / "pytorch_model.bin.index.json"
                if safe_index.exists() or bin_index.exists():
                    index_path = (
                        safe_index if safe_index.exists() else bin_index
                    )
                    with index_path.open("r", encoding="utf-8") as handle:
                        shard_names = sorted(
                            set(json.load(handle)["weight_map"].values())
                        )
                    state: dict[str, torch.Tensor] = {}
                    for shard_name in shard_names:
                        shard = Cola2ForConditionalGeneration._read_checkpoint_state(
                            path / shard_name, map_location=map_location
                        )
                        overlap = state.keys() & shard.keys()
                        if overlap:
                            raise RuntimeError(
                                f"Duplicate checkpoint keys: {sorted(overlap)[:5]}"
                            )
                        state.update(shard)
                    return state
                candidates = (
                    path / "model.safetensors",
                    path / "pytorch_model.bin",
                )
                path = next(
                    (candidate for candidate in candidates if candidate.exists()),
                    path,
                )
                if path.is_dir():
                    raise FileNotFoundError(
                        f"No model.safetensors or pytorch_model.bin in {path}."
                    )
            if path.suffix == ".safetensors":
                from safetensors.torch import load_file

                raw = load_file(str(path), device=str(map_location))
            else:
                raw = torch.load(
                    path, map_location=map_location, weights_only=True
                )
        if isinstance(raw, Mapping):
            for key in ("state_dict", "model"):
                nested = raw.get(key)
                if isinstance(nested, Mapping) and all(
                    torch.is_tensor(value) for value in nested.values()
                ):
                    raw = nested
                    break
        if not isinstance(raw, Mapping) or not all(
            isinstance(key, str) and torch.is_tensor(value)
            for key, value in raw.items()
        ):
            raise TypeError("Checkpoint must resolve to a tensor state dictionary.")
        return dict(raw)

    def load_strict_state_dict(
        self,
        source: Union[str, Path, Mapping[str, torch.Tensor]],
        *,
        map_location: Union[str, torch.device] = "cpu",
        prefix: Optional[str] = None,
        assign: bool = False,
    ) -> "Cola2ForConditionalGeneration":
        """Load package weights with exact key and shape validation.

        ``module.`` and ``_orig_mod.`` wrappers are stripped only when they wrap
        every key.  Other conversions must pass an explicit ``prefix`` so a
        partial or accidentally wrong native checkpoint cannot load silently.
        """

        state = self._read_checkpoint_state(source, map_location)
        if prefix is not None:
            if not all(key.startswith(prefix) for key in state):
                raise RuntimeError(
                    f"Not every checkpoint key starts with prefix {prefix!r}."
                )
            state = {key[len(prefix) :]: value for key, value in state.items()}
        for wrapper in ("module.", "_orig_mod."):
            if state and all(key.startswith(wrapper) for key in state):
                state = {
                    key[len(wrapper) :]: value for key, value in state.items()
                }

        expected = self.state_dict()
        missing = sorted(expected.keys() - state.keys())
        unexpected = sorted(state.keys() - expected.keys())
        mismatched = sorted(
            (
                key,
                tuple(state[key].shape),
                tuple(expected[key].shape),
            )
            for key in expected.keys() & state.keys()
            if tuple(state[key].shape) != tuple(expected[key].shape)
        )
        if missing or unexpected or mismatched:
            parts = []
            if missing:
                parts.append(f"missing={missing[:20]}")
            if unexpected:
                parts.append(f"unexpected={unexpected[:20]}")
            if mismatched:
                parts.append(f"shape_mismatches={mismatched[:20]}")
            raise RuntimeError(
                "Strict Cola2 state_dict validation failed: " + "; ".join(parts)
            )
        self.load_state_dict(state, strict=True, assign=assign)
        return self


def build_export_components(
    exported_config: Mapping[str, Any],
) -> dict[str, nn.Module]:
    """Build strict-validation components from a legacy export config.

    This compatibility factory accepts either a public Cola2 config or an
    earlier nested, path-free config and constructs the same DiT and VAE
    modules.
    """

    raw = dict(exported_config)
    if "dit_config" not in raw or "vae_config" not in raw:
        config = Cola2Config.from_dict(raw)
        return {"dit": Cola2DiT(config), "vae": Cola2CompressionVAE(config)}

    dit = dict(raw["dit_config"])
    vae = dict(raw["vae_config"])
    backbone = dict(vae.get("qwen3_backbone_config") or {})
    backbone = dict(backbone.get("text_config", backbone))
    model_config = Cola2Config(
        vocab_size=int(vae["vocab_size"]),
        encoder_embedding_size=(
            int(backbone.get("hidden_size", vae["dim"]))
            if bool(vae.get("qwen_encoder_token_weights", False))
            else int(vae["dim"])
        ),
        encoder_hidden_size=int(vae["dim"]),
        encoder_intermediate_size=int(vae["ffn_dim"]),
        encoder_num_attention_heads=int(vae["num_heads"]),
        encoder_num_hidden_layers=int(vae["token_encoder_num_blocks"]),
        pool_num_hidden_layers=int(vae["pool_num_blocks"]),
        latent_encoder_num_hidden_layers=int(
            vae["latent_encoder_num_blocks"]
        ),
        latent_decoder_num_hidden_layers=int(
            vae["latent_decoder_num_blocks"]
        ),
        latent_dim=int(vae["latent_dim"]),
        use_variational_encoder=bool(vae["use_variation"]),
        boundary_score_alpha=float(vae["scorer_score_alpha"]),
        boundary_score_threshold=float(vae["scorer_score_tau0"]),
        target_compression_rate=float(vae["rho_target"]),
        min_span_length=int(vae["lmin"]),
        max_span_length=int(vae["lmax"]),
        latent_block_causal=bool(vae["latent_block_causal"]),
        latent_block_size=int(vae["latent_block_size"]),
        dit_hidden_size=int(dit["txt_dim"]),
        dit_intermediate_size=int(dit["txt_dim"])
        * int(dit["expand_ratio"]),
        dit_num_hidden_layers=int(dit["num_layers"]),
        dit_num_attention_heads=int(dit["heads"]),
        dit_head_dim=int(dit["head_dim"]),
        dit_rope_dim=int(
            dit["rope_dim"][0]
            if isinstance(dit["rope_dim"], list)
            else dit["rope_dim"]
        ),
        dit_block_size=int(dit["block_size"]),
        decoder_hidden_size=int(
            backbone.get("hidden_size", vae["decoder_dim"])
        ),
        decoder_intermediate_size=int(
            backbone.get(
                "intermediate_size",
                int(vae["decoder_ffn_dim"]) // 2,
            )
        ),
        decoder_num_hidden_layers=int(vae["token_decoder_num_blocks"]),
        decoder_num_attention_heads=int(
            backbone.get(
                "num_attention_heads", vae["decoder_num_heads"]
            )
        ),
        decoder_num_key_value_heads=int(
            backbone.get(
                "num_key_value_heads", vae["decoder_num_heads"]
            )
        ),
        decoder_head_dim=int(
            backbone.get(
                "head_dim",
                int(vae["decoder_dim"])
                // int(vae["decoder_num_heads"]),
            )
        ),
        condition_num_attention_heads=int(vae["decoder_num_heads"]),
        decoder_condition_num_hidden_layers=int(vae["decoder_cond_blocks"]),
        decoder_lm_head_bias=bool(vae.get("bias", True)),
        duration_hidden_size=(
            None
            if vae.get("duration_hidden_dim") is None
            else int(vae["duration_hidden_dim"])
        ),
        local_position_dim=int(vae["loc_dim"]),
        hidden_dropout=float(vae["dropout"]),
        attention_dropout=float(vae["attn_dropout"]),
        layer_norm_eps=float(vae["layer_norm_eps"]),
        dit_layer_norm_eps=float(dit["norm_eps"]),
        encoder_rope_theta=float(vae["rope_theta"]),
        condition_rope_theta=float(vae["rope_theta"]),
        decoder_rope_theta=float(
            backbone.get("rope_theta", 1_000_000.0)
        ),
        dit_rope_theta=10_000.0,
        rope_theta=float(backbone.get("rope_theta", 1_000_000.0)),
        initializer_range=float(vae["init_std"]),
        scaling_factor=float(raw.get("latent_scaling_factor", 1.0)),
        shifting_factor=float(raw.get("latent_shifting_factor", 0.0)),
        pad_token_id=int(raw["pad_token_id"]),
        mask_token_id=int(raw["mask_token_id"]),
        eos_token_id=int(raw["eos_token_id"]),
        bos_token_id=int(raw["eos_token_id"]),
    )
    return {
        "dit": Cola2DiT(model_config),
        "vae": Cola2CompressionVAE(model_config),
    }