ymodel31_eval.py

"""Standalone evaluation/inference implementation for ymodel31.

This file intentionally contains a self-contained inference path so exported
checkpoints can be loaded without importing the training implementation.
Training-only features such as gradient checkpointing and self-distillation are
omitted here on purpose.
"""

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 safetensors.torch import load_file as load_safetensors
from transformers import GenerationMixin, PreTrainedModel
from transformers.activations import ACT2FN
from transformers.configuration_utils import PretrainedConfig
from transformers.modeling_outputs import CausalLMOutputWithPast


def normalize_gradient_checkpointing_level(value: Union[bool, int, str, None]) -> int:
    if isinstance(value, bool):
        return 1 if value else 0
    if value is None:
        return 0
    if isinstance(value, int):
        return max(0, value)
    text = str(value).strip().lower()
    if text in {"", "false", "off", "no", "none"}:
        return 0
    if text in {"true", "on", "yes"}:
        return 1
    try:
        return max(0, int(text))
    except ValueError as exc:
        raise ValueError(f"Unsupported gradient_checkpointing level: {value!r}") from exc


class YConfig31(PretrainedConfig):
    model_type = "ynet31"

    def __init__(
        self,
        dropout: float = 0.0,
        bos_token_id: int = 151644,
        eos_token_id: int = 151645,
        pad_token_id: int = 151643,
        hidden_act: str = "silu",
        hidden_size: int = 768,
        num_hidden_layers: int = 8,
        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",
        self_distill: bool = True,
        intermediate_size: int = 1536,
        num_heads: int = 12,
        mla_kv_lora_rank: int = 64,
        mla_qk_nope_head_dim: int = 64,
        mla_qk_rope_head_dim: int = 32,
        mla_attn_impl: str = "absorb",
        qkv_lora: bool = False,
        gradient_checkpointing: Union[bool, int, str] = 0,
        use_sengram: bool = True,
        sengram_bucket_size: Optional[int] = 4096,
        sengram_topk: int = 2,
        engram_bucket_size: Optional[int] = None,
        engram_topk: Optional[int] = None,
        **kwargs,
    ):
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            **kwargs,
        )
        self.dropout = dropout
        self.hidden_act = hidden_act
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.max_position_embeddings = max_position_embeddings
        self.vocab_size = vocab_size
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.rope_scaling = rope_scaling
        self.dtype = dtype
        self.self_distill = self_distill
        self.intermediate_size = intermediate_size
        self.num_heads = num_heads
        self.mla_kv_lora_rank = mla_kv_lora_rank
        self.mla_qk_nope_head_dim = mla_qk_nope_head_dim
        self.mla_qk_rope_head_dim = mla_qk_rope_head_dim
        self.mla_attn_impl = mla_attn_impl
        self.qkv_lora = qkv_lora
        self.gradient_checkpointing = normalize_gradient_checkpointing_level(gradient_checkpointing)
        self.use_sengram = bool(use_sengram)
        if engram_bucket_size is not None:
            sengram_bucket_size = engram_bucket_size
        if engram_topk is not None:
            sengram_topk = engram_topk
        self.sengram_bucket_size = sengram_bucket_size
        self.sengram_topk = sengram_topk
        self.engram_bucket_size = self.sengram_bucket_size
        self.engram_topk = self.sengram_topk

    @property
    def head_dim(self) -> int:
        return self.mla_qk_nope_head_dim + self.mla_qk_rope_head_dim

    @property
    def qk_head_dim(self) -> int:
        return self.head_dim

    def scale_lvl(self, lvl: int = 0):
        if lvl == 0:
            self.hidden_size = 768
            self.num_hidden_layers = 12
            self.num_heads = 8
            self.mla_kv_lora_rank = 256
            self.mla_qk_nope_head_dim = 128
            self.mla_qk_rope_head_dim = 64
            self.intermediate_size = 2048
            self.use_sengram = True
            self.sengram_bucket_size = 8192
            self.sengram_topk = 8
        elif lvl == -1:
            self.hidden_size = 768
            self.num_hidden_layers = 8
            self.num_heads = 6
            self.mla_kv_lora_rank = 128
            self.mla_qk_nope_head_dim = 64
            self.mla_qk_rope_head_dim = 64
            self.intermediate_size = 1536
            self.use_sengram = True
        elif lvl == -2:
            self.hidden_size = 512
            self.num_hidden_layers = 4
            self.num_heads = 4
            self.mla_kv_lora_rank = 128
            self.mla_qk_nope_head_dim = 64
            self.mla_qk_rope_head_dim = 64
            self.intermediate_size = 1024
            self.use_sengram = True
        else:
            raise ValueError(f"invalid ymodel31 scale level: {lvl}")
        return self


def _yarn_linear_ramp(low: float, high: float, dim: int) -> torch.Tensor:
    if low == high:
        high += 0.001
    linear = (torch.arange(dim, dtype=torch.float32) - low) / (high - low)
    return torch.clamp(linear, 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 precompute_freqs_cis(
    dim: int,
    end: int,
    theta: float,
    rope_scaling: Optional[dict] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    freqs = 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", end))
            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)
            freqs = freqs / factor * (1.0 - ramp) + freqs * ramp
            attention_factor = float(rope_scaling.get("attention_factor", 1.0))
    t = torch.arange(end)
    freqs = torch.outer(t, freqs).float()
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attention_factor
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attention_factor
    return freqs_cos, freqs_sin


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 = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
    elif cos.dim() == 3:
        cos = cos.unsqueeze(1)
        sin = sin.unsqueeze(1)
    return (x * cos) + (rotate_half(x) * sin)


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return (out * self.weight.float()).to(x.dtype)


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

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


class MLGA(nn.Module):
    """Multihead Latent Gated Attention"""

    def __init__(self, config: YConfig31, layer_id: int):
        super().__init__()
        self.layer_id = layer_id
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_heads
        self.dropout = config.dropout
        self.kv_lora_rank = config.mla_kv_lora_rank
        self.qk_nope_head_dim = config.mla_qk_nope_head_dim
        self.qk_rope_head_dim = config.mla_qk_rope_head_dim
        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
        self.attn_impl = config.mla_attn_impl
        self.softmax_scale = self.qk_head_dim**-0.5
        self.out_dim = self.num_heads * self.kv_lora_rank

        self.wq = nn.Linear(self.hidden_size, self.num_heads * self.qk_head_dim, bias=False)
        self.wkv_a = nn.Linear(self.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False)
        self.kv_norm = RMSNorm(self.kv_lora_rank, config.rms_norm_eps)
        self.wkv_b = nn.Linear(self.kv_lora_rank, self.num_heads * self.qk_nope_head_dim, bias=False)
        self.z_proj = nn.Linear(self.hidden_size, self.out_dim, bias=False)
        self.o_proj = nn.Linear(self.out_dim, self.hidden_size, bias=False)

    def _project_q(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        bsz, seq_len, _ = x.shape
        q = self.wq(x).reshape(bsz, seq_len, self.num_heads, self.qk_head_dim)
        return q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

    def _project_kv(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        raw = self.wkv_a(x)
        c_kv, k_pe = raw.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        c_kv = self.kv_norm(c_kv)
        k_pe = apply_rope_to_single(k_pe.unsqueeze(1), cos, sin).permute(0, 2, 1, 3)
        return c_kv, k_pe

    def _explicit_kv(self, c_kv: torch.Tensor, k_pe: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        bsz, seq_len, _ = c_kv.shape
        k_nope = self.wkv_b(c_kv).reshape(bsz, seq_len, self.num_heads, self.qk_nope_head_dim)
        k = torch.cat([k_nope, k_pe.expand(-1, -1, self.num_heads, -1)], dim=-1)
        v = c_kv.unsqueeze(2).expand(-1, -1, self.num_heads, -1)
        return k, v

    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:]
        mask = attention_mask.reshape(bsz, 1, 1, total_len).bool()
        return mask.expand(bsz, self.num_heads, seq_len, total_len)

    def _forward_sdpa(
        self,
        q_nope: torch.Tensor,
        q_pe: torch.Tensor,
        c_kv: torch.Tensor,
        k_pe: torch.Tensor,
        z: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        bsz, seq_len, _, _ = q_nope.shape
        total_len = c_kv.shape[1]
        k, v = self._explicit_kv(c_kv, k_pe)
        q = torch.cat([q_nope, q_pe], dim=-1).permute(0, 2, 1, 3)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        attn_mask = self._attention_mask(attention_mask, bsz, seq_len, total_len)
        is_causal = attention_mask is None and seq_len == total_len
        out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attn_mask,
            dropout_p=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 = out * torch.sigmoid(z)
        return self.o_proj(out)

    def _forward_absorb(
        self,
        q_nope: torch.Tensor,
        q_pe: torch.Tensor,
        c_kv: torch.Tensor,
        k_pe: torch.Tensor,
        z: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        bsz, seq_len, _, _ = q_nope.shape
        total_len = c_kv.shape[1]
        w = self.wkv_b.weight.reshape(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank)
        q_nope_c = torch.einsum("bshd,hdc->bshc", q_nope, w)
        scores = torch.einsum("bshc,btc->bsht", q_nope_c, c_kv)
        scores = scores + torch.einsum("bshr,btr->bsht", q_pe, k_pe.squeeze(2))
        scores = scores * self.softmax_scale

        causal = torch.full((seq_len, seq_len), float("-inf"), device=scores.device, dtype=scores.dtype)
        causal = torch.triu(causal, diagonal=1).reshape(1, seq_len, 1, seq_len)
        scores = scores + F.pad(causal, (total_len - seq_len, 0), value=0.0)
        if attention_mask is not None:
            if attention_mask.shape[-1] != total_len:
                attention_mask = attention_mask[..., -total_len:]
            scores = scores + (1.0 - attention_mask.reshape(bsz, 1, 1, total_len).float()) * -1e9
        probs = torch.softmax(scores.float(), dim=-1).to(q_nope.dtype)
        out = torch.einsum("bsht,btc->bshc", probs, c_kv).reshape(bsz, seq_len, self.out_dim)
        out = out * torch.sigmoid(z)
        return self.o_proj(out)

    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]]]:
        bsz, seq_len, _ = x.shape
        cos, sin = position_embeddings
        if cos.dim() == 2:
            cos = cos[:seq_len, : self.qk_rope_head_dim]
            sin = sin[:seq_len, : self.qk_rope_head_dim]
        else:
            cos = cos[:, :seq_len, : self.qk_rope_head_dim]
            sin = sin[:, :seq_len, : self.qk_rope_head_dim]
        q_nope, q_pe = self._project_q(x)
        q_pe = apply_rope_to_single(q_pe.permute(0, 2, 1, 3), cos, sin).permute(0, 2, 1, 3)
        c_kv, k_pe = self._project_kv(x, cos, sin)
        z = self.z_proj(x)

        if past_key_values is not None:
            past_c, past_pe = past_key_values
            c_kv = torch.cat([past_c, c_kv], dim=1)
            k_pe = torch.cat([past_pe, k_pe], dim=1)
        new_past = (c_kv, k_pe) if use_cache else None

        if self.attn_impl == "naive":
            out = self._forward_sdpa(q_nope, q_pe, c_kv, k_pe, z, attention_mask)
        else:
            out = self._forward_absorb(q_nope, q_pe, c_kv, k_pe, z, attention_mask)
        return out, new_past


class SwiGLU(nn.Module):
    def __init__(self, config: YConfig31, 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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        up, gate = self.up_proj(x), self.gate_proj(x)
        up = nn.functional.silu(gate) * up
        return self.down_proj(up)


class SengramIndexer(nn.Module):
    def __init__(self, config: YConfig31):
        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]:
        bucket_logits = self.bucket_proj(hidden_states)
        route_scores = torch.softmax(bucket_logits.float(), dim=-1)
        topk_ids = torch.topk(route_scores, k=self.topk, dim=-1, sorted=False).indices
        topk_scores = route_scores.gather(-1, topk_ids)
        denom = topk_scores.sum(dim=-1, keepdim=True).clamp_min(1e-20)
        topk_scores = (topk_scores / denom).to(bucket_logits.dtype)
        return topk_ids, topk_scores


class SengramPLE(nn.Module):
    def __init__(self, config: YConfig31):
        super().__init__()
        self.hidden_size = int(config.hidden_size)
        self.embedding = nn.Embedding(int(config.sengram_bucket_size or 4096), self.hidden_size)
        self.key_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.memory_norm = RMSNorm(self.hidden_size, config.rms_norm_eps)
        self.key_norm = RMSNorm(self.hidden_size, config.rms_norm_eps)
        self.query_norm = RMSNorm(self.hidden_size, config.rms_norm_eps)

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


class YBlock31(nn.Module):
    def __init__(self, config: YConfig31, 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.attn = MLGA(config, layer_id)
        self.ffn = SwiGLU(config)
        self.se1 = SEBlock(config.hidden_size, act=ACT2FN[config.hidden_act])
        self.se2 = SEBlock(config.hidden_size, act=ACT2FN[config.hidden_act])

    def forward(
        self,
        x: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        past_key_values: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        attention_mask: Optional[torch.Tensor] = None,
        route_ids: Optional[torch.Tensor] = None,
        route_scores: Optional[torch.Tensor] = None,
        **kwargs,
    ):
        if self.use_sengram and route_ids is not None and route_scores is not None and self.sengram_ple is not None:
            x = x + self.sengram_ple(x, route_ids, route_scores)
        x0 = self.se1(self.input_layernorm(x))
        attn_out, past = self.attn(
            x0,
            position_embeddings,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            use_cache=use_cache,
        )
        x = x + attn_out
        x0 = self.se2(self.post_attention_layernorm(x))
        x = x + self.ffn(x0)
        return x, past


class YModel31(nn.Module):
    def __init__(self, config: YConfig31):
        super().__init__()
        self.config = config
        self.vocab_size = config.vocab_size
        self.num_layers = config.num_hidden_layers
        self.dropout = config.dropout
        self.use_sengram = bool(config.use_sengram)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.sengram_indexer = SengramIndexer(config) if self.use_sengram else None
        self.layers = nn.ModuleList([YBlock31(config, i) for i in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        freqs_cos, freqs_sin = precompute_freqs_cis(
            dim=config.mla_qk_rope_head_dim,
            end=config.max_position_embeddings,
            theta=config.rope_theta,
            rope_scaling=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 _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        for key in list(state_dict.keys()):
            if key.startswith(prefix + "sengram."):
                state_dict.pop(key)
        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[list] = None,
        use_cache: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        **kwargs,
    ):
        bsz, seq_len = input_ids.shape
        if use_cache and past_key_values is None:
            past_key_values = [None] * self.num_layers
        if cache_position is None:
            if past_key_values is not None and past_key_values[0] is not None:
                past_seen = past_key_values[0][0].shape[1]
            else:
                past_seen = 0
            cache_position = torch.arange(past_seen, past_seen + seq_len, device=input_ids.device)

        x = self.embed_tokens(input_ids)
        if position_ids is None:
            position_ids = cache_position
        position_embeddings = (self.freqs_cos[position_ids].to(x.device), self.freqs_sin[position_ids].to(x.device))
        route_ids = None
        route_scores = None
        if self.use_sengram and self.sengram_indexer is not None:
            route_ids, route_scores = self.sengram_indexer(x)
        new_past = [] if use_cache else None

        for i, layer in enumerate(self.layers):
            past = past_key_values[i] if past_key_values is not None else None
            x, layer_past = layer(
                x,
                position_embeddings=position_embeddings,
                past_key_values=past,
                attention_mask=attention_mask,
                use_cache=use_cache,
                route_ids=route_ids,
                route_scores=route_scores,
            )
            if use_cache:
                new_past.append(layer_past)
        return self.norm(x), new_past


class YForCausalLM31(PreTrainedModel, GenerationMixin):
    config_class = YConfig31

    def __init__(self, config: Optional[YConfig31] = None):
        self.config = config or YConfig31()
        super().__init__(self.config)
        self.model = YModel31(self.config)
        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
        self.model.embed_tokens.weight = self.lm_head.weight
        self.OUT = CausalLMOutputWithPast()
        dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}.get(self.config.dtype)
        if dtype is not None:
            self.to(dtype)

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[list] = None,
        use_cache: bool = False,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ):
        h, past_kvs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            position_ids=kwargs.get("position_ids", None),
        )
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(h[:, slice_indices, :])
        self.OUT.__setitem__("last_hidden_state", h)
        self.OUT.__setitem__("logits", logits)
        self.OUT.__setitem__("past_key_values", past_kvs)
        return self.OUT

    def generate(
        self,
        inputs,
        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=True,
        num_return_sequences=1,
        do_sample=True,
        repetition_penalty=1.0,
        **kwargs,
    ):
        input_ids = kwargs.get("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", None)
        past_key_values = None
        if streamer:
            streamer.put(input_ids.cpu())
        with torch.no_grad():
            for _ in range(max_new_tokens):
                if use_cache and past_key_values is not None:
                    outputs = self.forward(input_ids[:, -1:], None, past_key_values, use_cache=use_cache)
                else:
                    outputs = self.forward(input_ids, attention_mask, past_key_values, use_cache=use_cache)
                logits = outputs.logits[:, -1, :] / temperature
                if repetition_penalty != 1.0:
                    for i in range(input_ids.shape[0]):
                        logits[i, torch.unique(input_ids[i])] /= repetition_penalty
                if logits_processor is not None:
                    logits = logits_processor(input_ids, logits)
                if top_k > 0:
                    logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float("inf")
                if top_p < 1.0:
                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                    mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
                    mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
                    logits[mask.scatter(1, sorted_indices, 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)
                past_key_values = outputs.past_key_values
                if streamer:
                    streamer.put(next_token.cpu())
                if eos_token_id and (next_token == eos_token_id).any():
                    break
        if streamer:
            streamer.end()
        return input_ids


def count_parameters(config: YConfig31) -> int:
    return sum(p.numel() for p in YForCausalLM31(config).parameters())


def _load_state_dict(path: Union[str, Path]) -> dict[str, torch.Tensor]:
    path = Path(path)
    if path.is_dir():
        safetensors_path = path / "model.safetensors"
        bin_path = path / "pytorch_model.bin"
        if safetensors_path.exists():
            path = safetensors_path
        elif bin_path.exists():
            path = bin_path
        else:
            raise FileNotFoundError(f"no model.safetensors or pytorch_model.bin found in {path}")
    if path.suffix == ".safetensors":
        return load_safetensors(str(path), device="cpu")
    return torch.load(path, map_location="cpu", weights_only=True)


def load_ymodel31_eval(path: Union[str, Path], config: Optional[YConfig31] = None, strict: bool = True) -> YForCausalLM31:
    path = Path(path)
    if config is None:
        config_path = path / "config.json" if path.is_dir() else path.with_name("config.json")
        if not config_path.exists():
            raise FileNotFoundError("config is required when config.json is not next to the checkpoint")
        config = YConfig31.from_json_file(str(config_path))
    model = YForCausalLM31(config)
    state = _load_state_dict(path)
    model.load_state_dict(state, strict=strict)
    model.eval()
    return model


YModel31Eval = YModel31
YForCausalLM31Eval = YForCausalLM31


__all__ = [
    "MLGA",
    "RMSNorm",
    "SEBlock",
    "SengramIndexer",
    "SengramPLE",
    "SwiGLU",
    "YBlock31",
    "YConfig31",
    "YForCausalLM31",
    "YForCausalLM31Eval",
    "YModel31",
    "YModel31Eval",
    "apply_rope_to_single",
    "count_parameters",
    "load_ymodel31_eval",
    "precompute_freqs_cis",
]