Blame for ymodel32_eval.py

"""Standalone inference implementation for YModel3.2.

This module contains the complete model definition.  Its only third-party
dependencies are PyTorch and Transformers.
"""

from __future__ import annotations

import math
from pathlib import Path
from typing import Optional, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GenerationMixin, PreTrainedModel
from transformers.configuration_utils import PretrainedConfig
from transformers.modeling_outputs import CausalLMOutputWithPast


def _normalize_group_sizes(group_size: Union[int, list[int], tuple[int, ...]]) -> tuple[int, ...]:
    if isinstance(group_size, bool):
        raise TypeError("group_size must be a positive integer or a list of integers")
    if isinstance(group_size, int):
        values = (group_size,)
    elif isinstance(group_size, (list, tuple)):
        values = tuple(group_size)
    else:
        raise TypeError("group_size must be a positive integer or a list of integers")
    if not values or any(
        isinstance(value, bool) or not isinstance(value, int) or value < 1
        for value in values
    ):
        raise ValueError("group_size must contain positive integers")
    if len(set(values)) != len(values):
        raise ValueError("group_size choices must be unique")
    return values


class YConfig32(PretrainedConfig):
    """Architecture config for YModel3.2."""

    model_type = "ynet32"
    MIN_FULL_BLOCKS = 3

    def __init__(
        self,
        dropout: float = 0.0,
        bos_token_id: int = 151644,
        eos_token_id: int = 151645,
        pad_token_id: int = 151643,
        hidden_size: int = 1024,
        num_hidden_layers: Optional[int] = None,
        max_position_embeddings: int = 8192,
        vocab_size: int = 6400,
        rms_norm_eps: float = 1e-6,
        rope_theta: float = 5e4,
        rope_scaling: Optional[dict] = None,
        dtype: str = "float32",
        intermediate_size: int = 2048,
        num_heads: int = 8,
        num_key_value_heads: Optional[int] = None,
        head_dim: int = 192,
        group_size: Union[int, list[int], tuple[int, ...]] = 4,
        layer_pattern: Optional[str] = "fgggfgggf",
        use_sengram: bool = False,
        sengram_bucket_size: Optional[int] = 4096,
        sengram_topk: int = 2,
        use_dcp: bool = False,
        dynamic_compression_probe_hidden_size: int = 768,
        tie_word_embeddings: bool = True,
        **kwargs,
    ):
        # Derived from group_size; must not be set independently.
        kwargs.pop("group_size_options", None)
        kwargs.pop("num_layers", None)
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )
        self.dropout = float(dropout)
        self.hidden_size = int(hidden_size)
        self.num_hidden_layers = int(num_hidden_layers) if num_hidden_layers is not None else 8
        self.max_position_embeddings = int(max_position_embeddings)
        self.vocab_size = int(vocab_size)
        self.rms_norm_eps = float(rms_norm_eps)
        self.rope_theta = float(rope_theta)
        self.rope_scaling = rope_scaling
        self.dtype = str(dtype)
        self.intermediate_size = int(intermediate_size)
        self.num_heads = int(num_heads)
        self.num_key_value_heads = (
            int(num_key_value_heads) if num_key_value_heads is not None else self.num_heads
        )
        self.head_dim = int(head_dim)
        self.group_size = list(group_size) if isinstance(group_size, tuple) else group_size
        self.group_size_options = list(_normalize_group_sizes(group_size))
        self.layer_pattern = layer_pattern
        self.use_sengram = bool(use_sengram)
        self.sengram_bucket_size = int(sengram_bucket_size or 4096)
        self.sengram_topk = int(sengram_topk)
        self.use_dcp = bool(use_dcp)
        self.dynamic_compression_probe_hidden_size = int(dynamic_compression_probe_hidden_size)
        if self.dynamic_compression_probe_hidden_size < 1:
            raise ValueError("dynamic_compression_probe_hidden_size must be positive")
        self.num_hidden_layers = len(resolve_layer_pattern(self))

    def scale_lvl(self, lvl: int = 0):
        """Apply a published YModel32 architecture preset."""
        if lvl == 2:
            self.hidden_size, self.num_heads, self.num_key_value_heads = 2048, 8, 2
            self.head_dim, self.intermediate_size = 256, 6144
            self.group_size, self.layer_pattern = 8, "fgggggfgggggfgggggfgggggf"
            self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 32768, 12
            self.use_dcp = False
        elif lvl == 1:
            self.hidden_size, self.num_heads, self.num_key_value_heads = 1024, 8, 2
            self.head_dim, self.intermediate_size = 256, 3072
            self.group_size, self.layer_pattern = 8, "fgggggfgggggfgggggfgggggfgggggf"
            self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 16384, 12
            self.use_dcp = False
            self.dynamic_compression_probe_hidden_size = 1024
        elif lvl == 0:
            self.hidden_size, self.num_heads, self.num_key_value_heads = 768, 12, 6
            self.head_dim, self.intermediate_size = 128, 2048
            self.group_size, self.layer_pattern = 8, "fggggfggggfggggfgggf"
            self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 8192, 12
            self.use_dcp = False
        elif lvl == -1:
            self.hidden_size, self.num_heads, self.num_key_value_heads = 768, 12, 6
            self.head_dim, self.intermediate_size = 128, 2048
            self.group_size, self.layer_pattern = list(range(4, 17)), "fggggfggggf"
            self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 4096, 12
            self.dynamic_compression_probe_hidden_size = 768
            self.use_dcp = True
        elif lvl == -2:
            self.hidden_size, self.num_heads, self.num_key_value_heads = 512, 4, 2
            self.head_dim, self.intermediate_size = 128, 1024
            self.group_size, self.layer_pattern = 8, "fggfggf"
            self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 1024, 8
            self.use_dcp = False
        else:
            raise ValueError(f"invalid YModel3.2 scale level: {lvl}")
        if self.layer_pattern.count("f") < self.MIN_FULL_BLOCKS:
            raise ValueError("YModel3.2 scale patterns require at least three full blocks")
        self.group_size_options = list(_normalize_group_sizes(self.group_size))
        self.num_hidden_layers = len(resolve_layer_pattern(self))
        return self


class RMSNorm(nn.RMSNorm):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__(dim, eps=eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return super().forward(x).to(x.dtype)


def _yarn_linear_ramp(low: float, high: float, dim: int) -> torch.Tensor:
    if low == high:
        high += 0.001
    return torch.clamp((torch.arange(dim, dtype=torch.float32) - low) / (high - low), 0.0, 1.0)


def _yarn_correction_dim(num_rotations: float, dim: int, theta: float, max_position_embeddings: int) -> float:
    return dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) / (2 * math.log(theta))


def resolve_rope_frequencies(
    dim: int,
    theta: float,
    rope_scaling: Optional[dict] = None,
    max_position_embeddings: Optional[int] = None,
) -> tuple[torch.Tensor, float]:
    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    attention_factor = 1.0
    if rope_scaling and str(rope_scaling.get("type", "yarn")).lower() == "yarn":
        factor = float(rope_scaling.get("factor", 1.0))
        if factor > 1.0:
            original = int(rope_scaling.get("original_max_position_embeddings", max_position_embeddings or 0))
            beta_fast = float(rope_scaling.get("beta_fast", 32.0))
            beta_slow = float(rope_scaling.get("beta_slow", 1.0))
            low = math.floor(_yarn_correction_dim(beta_fast, dim, theta, original))
            high = math.ceil(_yarn_correction_dim(beta_slow, dim, theta, original))
            ramp = _yarn_linear_ramp(low, high, dim // 2)
            inv_freq = inv_freq / factor * (1.0 - ramp) + inv_freq * ramp
            attention_factor = float(rope_scaling.get("attention_factor", 1.0))
    return inv_freq, attention_factor


def precompute_freqs_cis(
    dim: int,
    end: int,
    theta: float,
    rope_scaling: Optional[dict] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    inv_freq, attention_factor = resolve_rope_frequencies(
        dim, theta, rope_scaling, max_position_embeddings=end
    )
    freqs = torch.outer(torch.arange(end, dtype=torch.float32), inv_freq).float()
    return (
        torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attention_factor,
        torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attention_factor,
    )


def rope_embeddings_from_positions(
    positions: torch.Tensor,
    dim: int,
    theta: float,
    rope_scaling: Optional[dict] = None,
    max_position_embeddings: Optional[int] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    inv_freq, attention_factor = resolve_rope_frequencies(
        dim, theta, rope_scaling, max_position_embeddings
    )
    positions = positions.to(dtype=torch.float32)
    freqs = torch.outer(positions.reshape(-1), inv_freq.to(positions.device)).float()
    return (
        torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attention_factor,
        torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attention_factor,
    )


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    return torch.cat((-x[..., x.shape[-1] // 2 :], x[..., : x.shape[-1] // 2]), dim=-1)


def apply_rope_to_single(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    if cos.dim() == 2:
        cos, sin = cos.unsqueeze(0).unsqueeze(0), sin.unsqueeze(0).unsqueeze(0)
    elif cos.dim() == 3:
        cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
    return (x * cos) + (rotate_half(x) * sin)


class GatedAttention(nn.Module):
    """YModel32 gated grouped-query attention with Q/K RMS normalization."""

    def __init__(self, config: YConfig32, layer_id: int):
        super().__init__()
        del layer_id
        self.hidden_size = config.hidden_size
        self.num_heads = int(config.num_heads)
        self.num_kv_heads = int(config.num_key_value_heads or self.num_heads)
        if self.num_heads % self.num_kv_heads:
            raise ValueError("num_heads must be divisible by num_key_value_heads")
        self.num_kv_groups = self.num_heads // self.num_kv_heads
        self.dropout = config.dropout
        self.head_dim = int(config.head_dim)
        self.softmax_scale = self.head_dim**-0.5
        self.out_dim = self.num_heads * self.head_dim
        self.q_proj = nn.Linear(self.hidden_size, self.out_dim, bias=False)
        self.gate_proj = nn.Linear(self.hidden_size, self.out_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.out_dim, self.hidden_size, bias=False)
        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        q_key, gate_key = prefix + "q_proj.weight", prefix + "gate_proj.weight"
        legacy_qg = state_dict.get(q_key)
        if legacy_qg is not None and legacy_qg.shape[0] == self.out_dim * 2 and gate_key not in state_dict:
            legacy_qg = legacy_qg.reshape(self.num_heads, 2, self.head_dim, self.hidden_size)
            state_dict[q_key] = legacy_qg[:, 0].reshape(self.out_dim, self.hidden_size)
            state_dict[gate_key] = legacy_qg[:, 1].reshape(self.out_dim, self.hidden_size)
        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )

    def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
        if self.num_kv_groups == 1:
            return x
        bsz, num_kv_heads, seq_len, head_dim = x.shape
        x = x[:, :, None, :, :].expand(bsz, num_kv_heads, self.num_kv_groups, seq_len, head_dim)
        return x.reshape(bsz, self.num_heads, seq_len, head_dim)

    def _attention_mask(self, attention_mask: Optional[torch.Tensor], bsz: int, seq_len: int, total_len: int):
        if attention_mask is None:
            return None
        if attention_mask.shape[-1] != total_len:
            attention_mask = attention_mask[..., -total_len:]
        return attention_mask.reshape(bsz, 1, 1, total_len).bool().expand(bsz, self.num_heads, seq_len, total_len)

    def forward(
        self,
        x: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        past_key_values: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
        use_cache: bool = False,
        **kwargs,
    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
        del kwargs
        bsz, seq_len, _ = x.shape
        cos, sin = position_embeddings
        if cos.dim() == 2:
            cos, sin = cos[:seq_len, : self.head_dim], sin[:seq_len, : self.head_dim]
        else:
            cos, sin = cos[:, :seq_len, : self.head_dim], sin[:, :seq_len, : self.head_dim]
        q = self.q_proj(x).reshape(bsz, seq_len, self.num_heads, self.head_dim)
        gate = self.gate_proj(x)
        q = self.q_norm(q).permute(0, 2, 1, 3)
        k = self.k_norm(
            self.k_proj(x).reshape(bsz, seq_len, self.num_kv_heads, self.head_dim)
        ).permute(0, 2, 1, 3)
        v = self.v_proj(x).reshape(bsz, seq_len, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3)
        q = apply_rope_to_single(q, cos, sin).contiguous()
        k = apply_rope_to_single(k, cos, sin).contiguous()
        v = v.contiguous()
        if past_key_values is not None:
            past_k, past_v = past_key_values
            k, v = torch.cat([past_k, k], dim=2), torch.cat([past_v, v], dim=2)
        new_past = (k, v) if use_cache else None
        total_len = k.shape[2]
        k, v = self._repeat_kv(k).contiguous(), self._repeat_kv(v).contiguous()
        attn_mask = self._attention_mask(attention_mask, bsz, seq_len, total_len)
        is_causal = attention_mask is None and seq_len == total_len
        if q.dtype != v.dtype:
            q = q.to(dtype=v.dtype)
        if k.dtype != v.dtype:
            k = k.to(dtype=v.dtype)
        out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attn_mask,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=is_causal,
            scale=self.softmax_scale,
        )
        out = out.permute(0, 2, 1, 3).reshape(bsz, seq_len, self.out_dim)
        out = self.o_proj(out * torch.sigmoid(gate) * 2.0)
        return F.dropout(out, p=self.dropout, training=self.training), new_past


def resolve_layer_pattern(config: YConfig32) -> str:
    pattern = getattr(config, "layer_pattern", None)
    if pattern is None:
        pattern = "f" + "g" * max(0, int(config.num_hidden_layers) - 1)
    pattern = "".join(ch for ch in str(pattern).lower() if ch in {"f", "g"})
    if not pattern or pattern[0] != "f":
        raise ValueError(f"YModel3.2 layer_pattern must start with a full layer, got {pattern!r}")
    return pattern


class SwiGLU(nn.Module):
    def __init__(self, config: YConfig32, intermediate_size: Optional[int] = None):
        super().__init__()
        inter = intermediate_size or config.intermediate_size
        self.up_proj = nn.Linear(config.hidden_size, inter, bias=False)
        self.gate_proj = nn.Linear(config.hidden_size, inter, bias=False)
        self.down_proj = nn.Linear(inter, config.hidden_size, bias=False)
        self.dropout = config.dropout

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        up, gate = self.up_proj(x), self.gate_proj(x)
        return self.down_proj(F.dropout(F.silu(gate) * up, p=self.dropout, training=self.training))


class SEBlock(nn.Module):
    def __init__(self, dim: int, reduction: int = 16):
        super().__init__()
        reduction = max(reduction, dim // reduction)
        self.se = nn.Sequential(
            nn.Linear(dim, reduction, bias=False),
            nn.SiLU(),
            nn.Linear(reduction, dim, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * (2.0 * self.se(x))


class SengramIndexer(nn.Module):
    """Inference routing only: choose bucket top-k and normalized scores."""

    def __init__(self, config: YConfig32):
        super().__init__()
        self.hidden_size = int(config.hidden_size)
        self.bucket_size = int(config.sengram_bucket_size or 4096)
        self.topk = max(1, min(int(config.sengram_topk), self.bucket_size))
        self.bucket_proj = nn.Linear(self.hidden_size, self.bucket_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        router_input = hidden_states.float()
        router_input = router_input * torch.rsqrt(router_input.square().mean(dim=-1, keepdim=True) + 1e-6)
        bucket_logits = self.bucket_proj(router_input.to(dtype=hidden_states.dtype))
        scores = torch.sigmoid(bucket_logits.float())
        _, topk_ids = torch.topk(scores, k=self.topk, dim=-1, sorted=False)
        topk_scores = scores.gather(-1, topk_ids)
        topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True).clamp_min(1e-20)
        return topk_ids, topk_scores.to(bucket_logits.dtype)


class SengramPLE(nn.Module):
    def __init__(self, config: YConfig32):
        super().__init__()
        self.hidden_size = int(config.hidden_size)
        self.embedding = nn.Embedding(int(config.sengram_bucket_size or 4096), self.hidden_size)

    def forward(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_scores: torch.Tensor) -> torch.Tensor:
        topk_embed = self.embedding(topk_ids)
        return (topk_embed * topk_scores.unsqueeze(-1)).sum(dim=-2)


class DynamicCompressionProbe(nn.Module):
    """Inference probe used only when DCP controls token-to-group assignment."""

    def __init__(self, config: YConfig32):
        super().__init__()
        hidden_size = int(config.hidden_size)
        probe_size = int(config.dynamic_compression_probe_hidden_size)
        self.norm = RMSNorm(hidden_size, config.rms_norm_eps)
        self.gate_proj = nn.Linear(hidden_size, probe_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, probe_size, bias=False)
        self.down_proj = nn.Linear(probe_size, 1, bias=True)

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        hidden = self.norm(hidden.detach())
        hidden = F.silu(self.gate_proj(hidden)) * self.up_proj(hidden)
        return self.down_proj(hidden).squeeze(-1)


class G2F(nn.Module):
    """Add shifted group memory directly to the full-token stream."""

    def __init__(self, config: YConfig32):
        super().__init__()
        del config

    def forward(self, hidden: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
        return hidden + memory


class F2G(nn.Module):
    """Compress full-token states into groups, optionally DCP-weighted."""

    def __init__(self, config: YConfig32):
        super().__init__()
        self.hidden_size = int(config.hidden_size)
        self.group_size = _normalize_group_sizes(config.group_size)[0]

    def _aligned_weights(self, importance_logits, tokens: int, device: torch.device):
        if importance_logits is None:
            return None
        logits = importance_logits.to(device=device)
        if logits.ndim != 2:
            raise ValueError("DCP importance logits must have shape [batch, tokens]")
        required = max(0, tokens - 1)
        logits = logits[:, :required]
        if logits.shape[1] < required:
            logits = F.pad(logits, (0, required - logits.shape[1]), value=0.0)
        return torch.sigmoid(logits.float())

    def _weighted_reduce(self, h, attention_mask, group_ids, num_groups, importance_logits=None):
        bsz, tokens, hidden = h.shape
        if group_ids.shape != (bsz, tokens):
            raise ValueError(f"group_ids shape {tuple(group_ids.shape)} does not match hidden shape {(bsz, tokens)}")
        if attention_mask.shape != (bsz, tokens):
            raise ValueError(f"attention_mask shape {tuple(attention_mask.shape)} does not match hidden shape {(bsz, tokens)}")
        rest_h, rest_ids = h[:, 1:], group_ids[:, 1:]
        rest_valid = attention_mask[:, 1:].bool().reshape(-1)
        offsets = torch.arange(bsz, device=h.device, dtype=torch.int64).unsqueeze(1) * int(num_groups)
        flat_ids = (rest_ids.to(torch.int64) + offsets).reshape(-1)
        flat_h = rest_h.reshape(-1, hidden)
        aligned_weights = self._aligned_weights(importance_logits, tokens, h.device)
        if aligned_weights is None:
            flat_weights = torch.ones((bsz, max(0, tokens - 1)), device=h.device, dtype=torch.float32).reshape(-1)
        else:
            flat_weights = aligned_weights.reshape(-1)
        flat_ids, flat_h, flat_weights = flat_ids[rest_valid], flat_h[rest_valid], flat_weights[rest_valid]
        weighted = h.new_zeros((bsz * int(num_groups), hidden))
        weighted.index_add_(0, flat_ids, flat_h * flat_weights.to(h.dtype).unsqueeze(-1))
        denominators = torch.zeros((bsz * int(num_groups), 1), device=h.device, dtype=torch.float32)
        denominators.index_add_(0, flat_ids, flat_weights.unsqueeze(-1))
        compressed = weighted / denominators.clamp_min(1e-12).to(h.dtype)
        compressed = torch.where(denominators.gt(0.0), compressed, torch.zeros_like(compressed))
        compressed = compressed.reshape(bsz, int(num_groups), hidden)
        first = h[:, :1] * attention_mask[:, :1].to(h.dtype).unsqueeze(-1)
        return torch.cat([first, compressed[:, 1:]], dim=1)

    def forward(
        self,
        h: torch.Tensor,
        attention_mask: torch.Tensor,
        group_ids: Optional[torch.Tensor] = None,
        num_groups: Optional[int] = None,
        group_size: Optional[int] = None,
        importance_logits: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if group_ids is not None:
            if num_groups is None:
                raise ValueError("indexed F2G requires num_groups")
            return self.forward_indexed(h, attention_mask, group_ids, num_groups, importance_logits)
        group_size = self.group_size if group_size is None else int(group_size)
        bsz, padded_len, _ = h.shape
        rest_len = max(0, padded_len - 1)
        num_groups = 1 + (rest_len + group_size - 1) // group_size
        positions = torch.arange(padded_len, device=h.device, dtype=torch.int64)
        fixed_ids = torch.zeros((bsz, padded_len), device=h.device, dtype=torch.int64)
        if padded_len > 1:
            fixed_ids[:, 1:] = (positions[1:] - 1) // group_size + 1
        return self._weighted_reduce(h, attention_mask, fixed_ids, num_groups, importance_logits)

    def forward_indexed(self, h, attention_mask, group_ids, num_groups, importance_logits=None):
        return self._weighted_reduce(h, attention_mask, group_ids, int(num_groups), importance_logits)


class FullBlock(nn.Module):
    def __init__(self, config: YConfig32, layer_id: int, has_group_context: bool):
        super().__init__()
        self.has_group_context = has_group_context
        self.group_size = _normalize_group_sizes(config.group_size)[0]
        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.g2f = G2F(config) if has_group_context else None
        self.self_attn = GatedAttention(config, layer_id)
        self.mlp = SwiGLU(config)

    @property
    def attn(self):
        return self.self_attn

    @property
    def ffn(self):
        return self.mlp

    def _add_prev_group_residual(self, h, g, group_attention_mask=None, token_group_ids=None, group_size=None):
        if token_group_ids is not None:
            previous = (token_group_ids - 1).clamp_min(0)
            memory = g.gather(1, previous.unsqueeze(-1).expand(-1, -1, g.shape[-1]))
            memory = memory * token_group_ids.gt(0).unsqueeze(-1).to(memory.dtype)
            return self.g2f(h, memory)
        group_size = self.group_size if group_size is None else int(group_size)
        bsz, padded_len, hidden = h.shape
        rest_len = max(0, padded_len - 1)
        rest_groups = (rest_len + group_size - 1) // group_size
        groups = 1 + rest_groups
        if g.shape[1] != groups or rest_groups == 0:
            return h
        grouped_h = h[:, 1:].reshape(bsz, rest_groups, group_size, hidden)
        memory = g[:, :-1, :].unsqueeze(2)
        if group_attention_mask is not None:
            memory = memory * group_attention_mask[:, :-1].to(g.dtype).reshape(bsz, rest_groups, 1, 1)
        fused = self.g2f(grouped_h, memory).reshape(bsz, rest_groups * group_size, hidden)
        return torch.cat([h[:, :1], fused[:, :rest_len]], dim=1)

    def forward(self, h, g, token_position_embeddings, group_attention_mask=None, token_group_ids=None, group_size=None):
        if self.has_group_context and g is not None:
            h = self._add_prev_group_residual(h, g, group_attention_mask, token_group_ids, group_size)
        attn_out, _ = self.self_attn(self.input_layernorm(h), token_position_embeddings)
        h = h + attn_out
        return h + self.mlp(self.post_attention_layernorm(h))


class GroupBlock(nn.Module):
    def __init__(self, config: YConfig32, layer_id: int):
        super().__init__()
        self.use_sengram = bool(config.use_sengram)
        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.sengram_ple = SengramPLE(config) if self.use_sengram else None
        self.self_attn = GatedAttention(config, layer_id)
        self.mlp = SwiGLU(config)
        self.se1 = SEBlock(config.hidden_size)
        self.se2 = SEBlock(config.hidden_size)

    @property
    def attn(self):
        return self.self_attn

    @property
    def ffn(self):
        return self.mlp

    def forward(self, g, group_routes, group_position_embeddings):
        if self.use_sengram and group_routes is not None and self.sengram_ple is not None:
            g = g + self.sengram_ple(g, group_routes[0], group_routes[1])
        g0 = self.se1(self.input_layernorm(g))
        attn_out, _ = self.self_attn(g0, group_position_embeddings)
        g = g + attn_out
        return g + self.mlp(self.se2(self.post_attention_layernorm(g)))


class YModel32(nn.Module):
    """YModel3.2 backbone."""

    def __init__(self, config: YConfig32):
        super().__init__()
        self.config = config
        self.group_size_options = _normalize_group_sizes(config.group_size)
        self.group_size = self.group_size_options[0]
        self.max_group_size = max(self.group_size_options)
        self.active_group_size = self.group_size
        self.layer_pattern = resolve_layer_pattern(config)
        self.num_layers = len(self.layer_pattern)
        self.dropout = config.dropout
        self.use_sengram = bool(config.use_sengram)
        self.use_dcp = bool(config.use_dcp)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.sengram_indexer = SengramIndexer(config) if self.use_sengram else None
        # Kept in the state dict even when DCP is disabled so checkpoint
        # loading never drops learned parameters.
        self.dynamic_compression_probe = DynamicCompressionProbe(config)

        f2g_count = 0
        in_g = False
        for kind in self.layer_pattern:
            if kind == "g" and not in_g:
                f2g_count += 1
                in_g = True
            elif kind == "f":
                in_g = False
        self.f2g_layers = nn.ModuleList([F2G(config) for _ in range(f2g_count)])
        self.f2g_for_full = []
        f2g_ptr = 0
        for layer_pos, kind in enumerate(self.layer_pattern):
            if kind != "f":
                continue
            has_group_run = layer_pos + 1 < len(self.layer_pattern) and self.layer_pattern[layer_pos + 1] == "g"
            self.f2g_for_full.append(f2g_ptr if has_group_run else -1)
            if has_group_run:
                f2g_ptr += 1

        layers = []
        seen_full = 0
        for layer_id, kind in enumerate(self.layer_pattern):
            if kind == "f":
                layers.append(FullBlock(config, layer_id, has_group_context=seen_full > 0))
                seen_full += 1
            else:
                layers.append(GroupBlock(config, layer_id))
        self.layers = nn.ModuleList(layers)
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        freqs_cos, freqs_sin = precompute_freqs_cis(
            config.head_dim,
            config.max_position_embeddings + self.max_group_size,
            config.rope_theta,
            config.rope_scaling,
        )
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

    @property
    def sengram(self):
        return self.sengram_indexer

    def _sample_group_size(self) -> int:
        # Deterministic at inference: uses the first configured size.
        return self.group_size_options[0]

    def _pad_to_groups(self, input_ids, attention_mask, group_size, pad_to_groups=True):
        bsz, seq_len = input_ids.shape
        rest_len = max(0, seq_len - 1)
        pad_len = (-rest_len) % group_size if pad_to_groups else 0
        if pad_len:
            pad_id = self.config.pad_token_id if self.config.pad_token_id < self.config.vocab_size else 0
            input_ids = torch.cat([input_ids, input_ids.new_full((bsz, pad_len), pad_id)], dim=1)
            if attention_mask is not None:
                attention_mask = torch.cat([attention_mask, attention_mask.new_zeros((bsz, pad_len))], dim=1)
        if attention_mask is None:
            pad_id = self.config.pad_token_id if self.config.pad_token_id < self.config.vocab_size else -1
            attention_mask = input_ids.ne(pad_id).long()
        return input_ids, attention_mask, seq_len

    def _group_count(self, padded_len: int, group_size: int) -> int:
        return 1 + max(0, padded_len - 1) // group_size

    def _group_attention_mask(self, attention_mask, group_size):
        bsz, padded_len = attention_mask.shape
        rest_len = max(0, padded_len - 1)
        rest_groups = (rest_len + group_size - 1) // group_size
        first = attention_mask[:, :1]
        if rest_groups == 0:
            return first.long()
        rest_mask = attention_mask[:, 1:]
        pad_len = rest_groups * group_size - rest_len
        if pad_len:
            rest_mask = torch.cat([rest_mask, rest_mask.new_zeros((bsz, pad_len))], dim=1)
        return torch.cat([first, rest_mask.reshape(bsz, rest_groups, group_size).any(dim=-1).long()], dim=1)

    def _group_positions(self, groups: int, device: torch.device, group_size: int):
        if groups <= 1:
            return torch.zeros((groups,), device=device, dtype=torch.float32)
        rest = 1.0 + (torch.arange(groups - 1, device=device, dtype=torch.float32) + 0.5) * group_size
        return torch.cat([rest.new_zeros((1,)), rest], dim=0)

    def _dcp_group_ids(self, probe_logits, attention_mask, group_size=None):
        bsz, tokens = probe_logits.shape
        group_ids = torch.zeros((bsz, tokens), device=probe_logits.device, dtype=torch.int64)
        if tokens <= 1:
            return group_ids, 1
        group_size = self.group_size if group_size is None else int(group_size)
        target_groups = max(1, (tokens + group_size - 1) // group_size)
        target_rest_groups = target_groups - 1
        if target_rest_groups == 0:
            return group_ids, target_groups
        score = torch.sigmoid(probe_logits[:, :-1]).float()
        score = score * attention_mask[:, 1:].to(score.dtype)
        budget = float(group_size) * 0.5
        raw_integral = score.sum(dim=1, keepdim=True) / budget
        scale = target_rest_groups / raw_integral.clamp_min(1e-6)
        scaled_score = score * scale
        integral_before_token = (torch.cumsum(scaled_score, dim=1) - scaled_score) / budget
        integral_bucket = torch.floor(integral_before_token).to(torch.int64)
        integral_bucket = integral_bucket.clamp_(0, target_rest_groups - 1)
        valid_rest = attention_mask[:, 1:].bool()
        rest_positions = torch.arange(tokens - 1, device=probe_logits.device).unsqueeze(0)
        last_valid = rest_positions.masked_fill(~valid_rest, -1).amax(dim=1)
        safe_last_valid = last_valid.clamp_min(0)
        rows = torch.arange(bsz, device=probe_logits.device)
        current_last_bucket = integral_bucket[rows, safe_last_valid]
        final_bucket = torch.where(
            last_valid.ge(0),
            torch.full_like(current_last_bucket, target_rest_groups - 1),
            current_last_bucket,
        )
        integral_bucket[rows, safe_last_valid] = final_bucket
        group_ids[:, 1:] = integral_bucket + 1
        return group_ids, target_groups

    def _indexed_group_metadata(self, group_ids, attention_mask, num_groups):
        bsz, tokens = group_ids.shape
        offsets = torch.arange(bsz, device=group_ids.device, dtype=torch.int64).unsqueeze(1) * num_groups
        flat_ids = (group_ids + offsets).reshape(-1)
        valid = attention_mask.reshape(-1).bool()
        counts = torch.zeros(bsz * num_groups, device=group_ids.device, dtype=torch.float32)
        counts.index_add_(0, flat_ids[valid], torch.ones_like(flat_ids[valid], dtype=torch.float32))
        positions = torch.arange(tokens, device=group_ids.device, dtype=torch.float32).expand(bsz, -1).reshape(-1)
        position_sums = torch.zeros_like(counts)
        position_sums.index_add_(0, flat_ids[valid], positions[valid])
        group_mask = counts.reshape(bsz, num_groups).gt(0)
        centers = (position_sums / counts.clamp_min(1.0)).reshape(bsz, num_groups)
        return group_mask.long(), centers

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> torch.Tensor:
        del kwargs
        if input_ids is None:
            raise ValueError("input_ids is required")
        group_size = self._sample_group_size()
        self.active_group_size = group_size
        input_ids, attention_mask, orig_len = self._pad_to_groups(
            input_ids, attention_mask, group_size, pad_to_groups=not self.use_dcp
        )
        bsz, padded_len = input_ids.shape
        h = F.dropout(self.embed_tokens(input_ids), p=self.dropout, training=False)
        if position_ids is None:
            position_ids = torch.arange(padded_len, device=input_ids.device)
        token_position_embeddings = (
            self.freqs_cos[position_ids[:padded_len]].to(h.device),
            self.freqs_sin[position_ids[:padded_len]].to(h.device),
        )
        groups = self._group_count(padded_len, group_size)
        group_attention_mask = self._group_attention_mask(attention_mask, group_size)
        group_center_positions = self._group_positions(groups, input_ids.device, group_size)
        group_position_embeddings = rope_embeddings_from_positions(
            group_center_positions,
            self.config.head_dim,
            self.config.rope_theta,
            self.config.rope_scaling,
            self.config.max_position_embeddings,
        )
        group_position_embeddings = (group_position_embeddings[0].to(h.device), group_position_embeddings[1].to(h.device))

        g = None
        pending_group_update = None
        group_routes = None
        token_group_ids = None
        dcp_num_groups = None
        full_idx = 0
        difficulty_logits = None
        for layer_pos, kind in enumerate(self.layer_pattern):
            if kind == "f":
                h = self.layers[layer_pos](
                    h,
                    g,
                    token_position_embeddings,
                    group_attention_mask=group_attention_mask,
                    token_group_ids=token_group_ids,
                    group_size=group_size,
                )
                if difficulty_logits is None and self.use_dcp:
                    difficulty_logits = self.dynamic_compression_probe(h)
                    token_group_ids, dcp_num_groups = self._dcp_group_ids(
                        difficulty_logits, attention_mask, group_size
                    )
                    group_attention_mask, group_center_positions = self._indexed_group_metadata(
                        token_group_ids, attention_mask, dcp_num_groups
                    )
                    group_position_embeddings = rope_embeddings_from_positions(
                        group_center_positions,
                        self.config.head_dim,
                        self.config.rope_theta,
                        self.config.rope_scaling,
                        self.config.max_position_embeddings,
                    )
                    group_position_embeddings = (
                        group_position_embeddings[0].reshape(bsz, dcp_num_groups, -1).to(h.device),
                        group_position_embeddings[1].reshape(bsz, dcp_num_groups, -1).to(h.device),
                    )
                    groups = dcp_num_groups
                f2g_idx = self.f2g_for_full[full_idx]
                if f2g_idx >= 0:
                    if token_group_ids is not None and dcp_num_groups is not None:
                        pending_group_update = self.f2g_layers[f2g_idx](
                            h,
                            attention_mask,
                            token_group_ids,
                            dcp_num_groups,
                            importance_logits=difficulty_logits,
                        )
                    else:
                        pending_group_update = self.f2g_layers[f2g_idx](
                            h,
                            attention_mask,
                            group_size=group_size,
                            importance_logits=difficulty_logits,
                        )
                if self.sengram_indexer is not None and pending_group_update is not None:
                    group_routes = self.sengram_indexer(pending_group_update)
                full_idx += 1
            else:
                if pending_group_update is not None:
                    g = pending_group_update if g is None else g + pending_group_update
                    pending_group_update = None
                elif g is None:
                    g = h.new_zeros(h.shape[0], groups, h.shape[-1])
                g = self.layers[layer_pos](g, group_routes, group_position_embeddings)
        return self.norm(h)[:, :orig_len]


class YForCausalLM32(PreTrainedModel, GenerationMixin):
    config_class = YConfig32
    base_model_prefix = "model"
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: Optional[YConfig32] = None):
        config = config or YConfig32()
        super().__init__(config)
        self.model = YModel32(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.model.embed_tokens.weight = self.lm_head.weight
        dtype_name = config.dtype
        if dtype_name not in {"float16", "bfloat16", "float32"}:
            hf_dtype = getattr(config, "torch_dtype", None)
            dtype_name = str(hf_dtype).replace("torch.", "") if hf_dtype is not None else "float32"
        dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}.get(dtype_name)
        if dtype is not None:
            self.to(dtype)

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        position_ids = kwargs.pop("position_ids", None)
        # This architecture currently recomputes the full prefix. Cache inputs
        # are accepted for API compatibility but are not consumed.
        kwargs.pop("past_key_values", None)
        kwargs.pop("use_cache", None)
        kwargs.pop("cache_position", None)
        hidden = self.model(input_ids, attention_mask, position_ids=position_ids)
        indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        output = CausalLMOutputWithPast(
            logits=self.lm_head(hidden[:, indices, :]),
            past_key_values=None,
        )
        output["last_hidden_state"] = hidden
        return output

    def generate(
        self,
        inputs=None,
        attention_mask=None,
        max_new_tokens=8192,
        temperature=0.85,
        top_p=0.85,
        top_k=50,
        eos_token_id=None,
        streamer=None,
        use_cache=False,
        num_return_sequences=1,
        do_sample=True,
        repetition_penalty=1.0,
        **kwargs,
    ):
        del use_cache
        if inputs is None:
            inputs = kwargs.pop("input_ids", None)
        if inputs is None:
            raise ValueError("inputs or input_ids must be provided")
        input_ids = inputs.repeat(num_return_sequences, 1)
        attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
        logits_processor = kwargs.get("logits_processor")
        eos_token_id = self.config.eos_token_id if eos_token_id is None else eos_token_id
        if streamer is not None:
            streamer.put(input_ids.cpu())
        with torch.inference_mode():
            for _ in range(max_new_tokens):
                outputs = self.forward(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    logits_to_keep=1,
                    use_cache=False,
                )
                logits = outputs.logits[:, -1, :]
                if temperature and temperature > 0:
                    logits = logits / temperature
                if repetition_penalty != 1.0:
                    for i in range(input_ids.shape[0]):
                        seen = torch.unique(input_ids[i])
                        logits[i, seen] = torch.where(
                            logits[i, seen] < 0,
                            logits[i, seen] * repetition_penalty,
                            logits[i, seen] / repetition_penalty,
                        )
                if logits_processor is not None:
                    logits = logits_processor(input_ids, logits)
                if top_k > 0:
                    k = min(int(top_k), logits.shape[-1])
                    cutoff = torch.topk(logits, k, dim=-1)[0][..., -1, None]
                    logits = logits.masked_fill(logits < cutoff, -float("inf"))
                if top_p < 1.0:
                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                    sorted_mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
                    sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
                    sorted_mask[..., 0] = False
                    remove_mask = torch.zeros_like(sorted_mask).scatter(1, sorted_indices, sorted_mask)
                    logits = logits.masked_fill(remove_mask, -float("inf"))
                next_token = (
                    torch.multinomial(torch.softmax(logits, dim=-1), 1)
                    if do_sample
                    else torch.argmax(logits, dim=-1, keepdim=True)
                )
                input_ids = torch.cat([input_ids, next_token], dim=-1)
                if attention_mask is not None:
                    attention_mask = torch.cat(
                        [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
                    )
                if streamer is not None:
                    streamer.put(next_token.cpu())
                if eos_token_id is not None and bool((next_token == eos_token_id).all()):
                    break
        if streamer is not None:
            streamer.end()
        return input_ids


def count_parameters(config: YConfig32) -> int:
    return sum(parameter.numel() for parameter in YForCausalLM32(config).parameters())


def _state_dict_from_checkpoint(payload) -> dict[str, torch.Tensor]:
    if not isinstance(payload, dict):
        state_dict = getattr(payload, "state_dict", None)
        if callable(state_dict):
            payload = state_dict()
        else:
            raise TypeError("checkpoint must contain a PyTorch state-dict mapping")
    for key in ("state_dict", "model_state_dict", "model", "student", "student_state_dict"):
        nested = payload.get(key)
        if isinstance(nested, dict) and nested and all(isinstance(k, str) for k in nested):
            if any(torch.is_tensor(value) for value in nested.values()):
                payload = nested
                break
    if not payload or not all(isinstance(key, str) for key in payload):
        raise TypeError("checkpoint does not look like a PyTorch state dict")
    return payload


def _load_state_dict(path: Union[str, Path]) -> dict[str, torch.Tensor]:
    path = Path(path)
    if path.is_dir():
        candidates = ("sft.pt", "distill.pt", "model.pt", "pytorch_model.bin")
        path = next((path / name for name in candidates if (path / name).is_file()), None)
        if path is None:
            raise FileNotFoundError("no .pt/.bin PyTorch checkpoint found in the directory")
    if path.suffix.lower() == ".safetensors":
        raise ValueError(".safetensors loading is not included; export a PyTorch .pt state dict")
    try:
        payload = torch.load(path, map_location="cpu", weights_only=True)
    except TypeError:
        payload = torch.load(path, map_location="cpu")
    return _state_dict_from_checkpoint(payload)


def _key_variants(key: str) -> tuple[str, ...]:
    variants = [key]
    current = key
    while current.startswith("module.") or current.startswith("_orig_mod."):
        current = current.split(".", 1)[1]
        variants.append(current)
    for source in tuple(variants):
        aliases = (
            source.replace(".attn.", ".self_attn."),
            source.replace(".ffn.", ".mlp."),
            source.replace(".ffn.dense.", ".mlp."),
            source.replace("model.sengram.", "model.sengram_indexer."),
        )
        variants.extend(aliases)
    # Some compiled checkpoints insert _orig_mod below the outer model prefix.
    variants.extend(
        key.replace("model._orig_mod.", "model.").replace("._orig_mod.", ".")
        for key in tuple(variants)
    )
    return tuple(dict.fromkeys(variants))


def _load_compatible_weights(model: YForCausalLM32, source: dict[str, torch.Tensor], strict: bool = True):
    target = model.state_dict()
    normalized = {}
    for key, value in source.items():
        if not torch.is_tensor(value):
            continue
        for variant in _key_variants(key):
            normalized.setdefault(variant, value)
    matched = {}
    shape_mismatches = []
    for key, tensor in target.items():
        found = next((normalized[name] for name in _key_variants(key) if name in normalized), None)
        if found is None:
            continue
        if tuple(found.shape) != tuple(tensor.shape):
            # Early checkpoints packed Q and gate projections into one
            # tensor; split them in the same head-major layout used by the
            # checkpoint loader.
            if (
                key.endswith(".q_proj.weight")
                and found.ndim == 2
                and found.shape[0] == tensor.shape[0] * 2
                and found.shape[1] == tensor.shape[1]
            ):
                packed = found.reshape(
                    model.config.num_heads,
                    2,
                    model.config.head_dim,
                    model.config.hidden_size,
                )
                matched[key] = packed[:, 0].reshape_as(tensor)
                gate_key = key[: -len("q_proj.weight")] + "gate_proj.weight"
                matched[gate_key] = packed[:, 1].reshape_as(target[gate_key])
                continue
            shape_mismatches.append((key, tuple(found.shape), tuple(tensor.shape)))
            continue
        matched[key] = found
    result = model.load_state_dict(matched, strict=False)
    missing = list(result.missing_keys)
    if model.config.tie_word_embeddings:
        tied = {"model.embed_tokens.weight", "lm_head.weight"}
        if tied.intersection(matched):
            missing = [key for key in missing if key not in tied]
    if shape_mismatches and strict:
        preview = "; ".join(f"{key}: checkpoint{src} != model{dst}" for key, src, dst in shape_mismatches[:8])
        raise RuntimeError(f"checkpoint tensor shape mismatch: {preview}")
    if missing and strict:
        preview = ", ".join(missing[:20])
        raise RuntimeError(f"checkpoint is missing {len(missing)} model tensors: {preview}")
    unique_target = {}
    for key, tensor in target.items():
        identity = (tensor.device.type, tensor.data_ptr(), tuple(tensor.shape))
        unique_target.setdefault(identity, []).append(key)
    loaded_unique = sum(any(key in matched for key in keys) for keys in unique_target.values())
    print(f"loaded: {loaded_unique}/{len(unique_target)} unique model tensors")
    if shape_mismatches:
        print(f"shape-mismatched tensors ignored: {len(shape_mismatches)}")
    if missing:
        print(f"missing model tensors: {len(missing)}")
    return model


def load_ymodel32_eval(
    path: Union[str, Path],
    config: Optional[YConfig32] = None,
    strict: bool = True,
    device: Union[str, torch.device] = "cpu",
) -> YForCausalLM32:
    """Load a standalone evaluator from a .pt/.bin state dict and config JSON."""
    path = Path(path)
    if config is None:
        config_candidates = (
            path / "config.json" if path.is_dir() else path.with_name("config.json"),
            path.parent.parent / "config.json" if not path.is_dir() else path.parent / "config.json",
        )
        config_path = next((candidate for candidate in config_candidates if candidate.is_file()), None)
        if config_path is None:
            raise FileNotFoundError("pass config=YConfig32(...) or place config.json next to the checkpoint")
        config = YConfig32.from_json_file(str(config_path))
    model = YForCausalLM32(config)
    model = _load_compatible_weights(model, _load_state_dict(path), strict=strict)
    return model.to(device).eval()


YModel32Eval = YModel32
YForCausalLM32Eval = YForCausalLM32


__all__ = [
    "DynamicCompressionProbe",
    "F2G",
    "FullBlock",
    "G2F",
    "GatedAttention",
    "GroupBlock",
    "RMSNorm",
    "SengramIndexer",
    "SengramPLE",
    "SwiGLU",
    "YConfig32",
    "YForCausalLM32",
    "YForCausalLM32Eval",
    "YModel32",
    "YModel32Eval",
    "apply_rope_to_single",
    "count_parameters",
    "load_ymodel32_eval",
    "precompute_freqs_cis",
    "resolve_layer_pattern",
    "resolve_rope_frequencies",
    "rope_embeddings_from_positions",
]