Blame for modeling_comet.py

"""
Comet-1 (OpenModel-1) for Hugging Face Transformers.

Inference/fine-tuning port of the ``OpenModel1`` module defined in the training
pipeline (``train.py``). Parameter names and shapes are identical to the training
checkpoint, so ``model.safetensors`` loads without any key remapping.

Architecture
------------
A pre-norm decoder with tied input/output embeddings. Blocks follow a repeating
schedule (``config.layer_pattern``, by default ``[global, global, global, recall]``):

* ``global`` blocks use causal grouped-query attention with per-head RMSNorm on
  queries and keys (QK-norm) followed by rotary position embeddings.
* ``recall`` blocks replace attention with a gated, depthwise, dilated *causal*
  1-D convolution over a narrower lane (``config.memory_width``).

Every block then applies a SwiGLU feed-forward. Both residual branches are scaled
by a learned per-channel vector (LayerScale).

Compared with ``train.py`` this file adds what generation needs and the training
loop did not: an incremental cache (attention keys/values *and* convolution
history), padding-aware attention masks, and the standard Transformers outputs.
"""

import math
from typing import Optional, Tuple, Union

import torch
import torch.nn.functional as F
from torch import nn

from transformers import DynamicCache, GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast

from .configuration_comet import CometConfig


def _cache_length(cache) -> int:
    """Number of tokens already stored in ``cache`` (0 for ``None`` or an empty cache)."""
    if cache is None:
        return 0
    try:
        return int(cache.get_seq_length())
    except Exception:
        return 0


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.weight.to(x.dtype)
        if hasattr(F, "rms_norm"):
            return F.rms_norm(x, (x.shape[-1],), weight, self.eps)
        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
        return weight * (x * torch.rsqrt(variance + self.eps).to(x.dtype))

    def extra_repr(self) -> str:
        return f"{self.weight.shape[0]}, eps={self.eps}"


def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Rotate interleaved (even, odd) channel pairs. ``x``: [B, H, T, D]; ``cos``/``sin``: [B|1, 1, T, D/2]."""
    even = x[..., ::2]
    odd = x[..., 1::2]
    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)


class CometAttention(nn.Module):
    """Causal grouped-query attention with QK-norm and RoPE."""

    def __init__(self, config: CometConfig):
        super().__init__()
        self.heads = config.num_attention_heads
        self.kv_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.q_dim = self.heads * self.head_dim
        self.kv_dim = self.kv_heads * self.head_dim
        self.fused_projections = config.fused_projections
        if self.fused_projections:
            self.qkv_proj = nn.Linear(config.hidden_size, self.q_dim + 2 * self.kv_dim, bias=False)
        else:
            self.q_proj = nn.Linear(config.hidden_size, self.q_dim, bias=False)
            self.k_proj = nn.Linear(config.hidden_size, self.kv_dim, bias=False)
            self.v_proj = nn.Linear(config.hidden_size, self.kv_dim, bias=False)
        self.o_proj = nn.Linear(self.q_dim, config.hidden_size, bias=False)
        self.o_proj._comet_residual_out = True
        if config.qk_norm:
            self.q_norm = CometRMSNorm(self.head_dim, eps=config.rms_norm_eps)
            self.k_norm = CometRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        else:
            self.q_norm = nn.Identity()
            self.k_norm = nn.Identity()

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None,
        is_causal: bool = False,
        cache=None,
        layer_idx: int = 0,
    ) -> torch.Tensor:
        batch, sequence, _ = x.shape
        if self.fused_projections:
            q, k, v = self.qkv_proj(x).split((self.q_dim, self.kv_dim, self.kv_dim), dim=-1)
        else:
            q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
        q = q.view(batch, sequence, self.heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, sequence, self.kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, sequence, self.kv_heads, self.head_dim).transpose(1, 2)

        cos = cos.to(device=q.device, dtype=q.dtype)
        sin = sin.to(device=q.device, dtype=q.dtype)
        q = _apply_rope(self.q_norm(q), cos, sin)
        k = _apply_rope(self.k_norm(k), cos, sin)

        if cache is not None:
            # Keys are stored already normalised and rotated, at key/value-head granularity.
            k, v = cache.update(k, v, layer_idx)

        groups = self.heads // self.kv_heads
        if groups > 1:
            # Query head h reads key/value head h // groups (same grouping as `enable_gqa`).
            length = k.shape[-2]
            k = k[:, :, None].expand(batch, self.kv_heads, groups, length, self.head_dim)
            k = k.reshape(batch, self.heads, length, self.head_dim)
            v = v[:, :, None].expand(batch, self.kv_heads, groups, length, self.head_dim)
            v = v.reshape(batch, self.heads, length, self.head_dim)

        if attn_mask is not None:
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask.to(q.device))
        else:
            y = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
        return self.o_proj(y.transpose(1, 2).reshape(batch, sequence, self.q_dim))


class CometRecallMixer(nn.Module):
    """Gated local recall branch: depthwise, dilated, causal convolution over a narrow lane."""

    def __init__(self, config: CometConfig, dilation: int):
        super().__init__()
        width = config.memory_width
        self.width = width
        self.kernel_size = config.memory_kernel_size
        self.dilation = dilation
        # Number of past frames a single output position can see.
        self.history = (self.kernel_size - 1) * self.dilation
        self.in_proj = nn.Linear(config.hidden_size, 2 * width, bias=False)
        self.depthwise = nn.Conv1d(
            width,
            width,
            kernel_size=self.kernel_size,
            dilation=dilation,
            groups=width,
            bias=False,
        )
        self.out_proj = nn.Linear(width, config.hidden_size, bias=False)
        self.out_proj._comet_residual_out = True

    def forward(
        self,
        x: torch.Tensor,
        token_mask: Optional[torch.Tensor] = None,
        cache=None,
        layer_idx: int = 0,
    ) -> torch.Tensor:
        batch, sequence, _ = x.shape
        gate, value = self.in_proj(x).chunk(2, dim=-1)
        if token_mask is not None:
            # Padding tokens contribute nothing to the convolution, exactly like the
            # zero padding that precedes the first real token during training.
            value = value * token_mask.to(device=value.device, dtype=value.dtype).unsqueeze(-1)

        if cache is not None:
            # The convolution input is cached in the "key" slot as [B, 1, T, width] so
            # that every generic cache operation (beam reordering, cropping, batch
            # expansion) applies to it unchanged. The "value" slot is a placeholder.
            frames, _ = cache.update(value.unsqueeze(1), value.new_zeros(batch, 1, sequence, 1), layer_idx)
            context = frames[:, 0, -(self.history + sequence) :]
        else:
            context = value

        left = self.history + sequence - context.shape[1]
        context = F.pad(context.transpose(1, 2), (left, 0))
        value = self.depthwise(context).transpose(1, 2)
        return self.out_proj(torch.sigmoid(gate) * value)


class CometSwiGLU(nn.Module):
    def __init__(self, config: CometConfig):
        super().__init__()
        self.fused_projections = config.fused_projections
        if self.fused_projections:
            self.gate_up = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
        else:
            self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
            self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.down._comet_residual_out = True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.fused_projections:
            gate, up = self.gate_up(x).chunk(2, dim=-1)
        else:
            gate, up = self.gate(x), self.up(x)
        return self.down(F.silu(gate) * up)


class CometBlock(nn.Module):
    def __init__(self, config: CometConfig, index: int):
        super().__init__()
        self.index = index
        self.kind = config.block_kind(index)
        self.norm1 = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.norm2 = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        if self.kind == "global":
            self.mixer = CometAttention(config)
        else:
            self.mixer = CometRecallMixer(config, config.block_dilation(index))
        self.mlp = CometSwiGLU(config)
        self.mixer_scale = nn.Parameter(torch.full((config.hidden_size,), float(config.layer_scale_init)))
        self.mlp_scale = nn.Parameter(torch.full((config.hidden_size,), float(config.layer_scale_init)))
        self.dropout = config.dropout

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None,
        token_mask: Optional[torch.Tensor] = None,
        is_causal: bool = False,
        cache=None,
    ) -> torch.Tensor:
        h = self.norm1(x)
        if self.kind == "global":
            h = self.mixer(h, cos, sin, attn_mask=attn_mask, is_causal=is_causal, cache=cache, layer_idx=self.index)
        else:
            h = self.mixer(h, token_mask=token_mask, cache=cache, layer_idx=self.index)
        if self.dropout:
            h = F.dropout(h, p=self.dropout, training=self.training)
        x = x + h * self.mixer_scale.to(h.dtype)
        h = self.mlp(self.norm2(x))
        if self.dropout:
            h = F.dropout(h, p=self.dropout, training=self.training)
        return x + h * self.mlp_scale.to(h.dtype)


class CometPreTrainedModel(PreTrainedModel):
    config_class = CometConfig
    base_model_prefix = ""
    supports_gradient_checkpointing = True
    _no_split_modules = ["CometBlock"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_cache_class = True

    def _init_weights(self, module: nn.Module) -> None:
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            if getattr(module, "_comet_residual_out", False):
                std = std / math.sqrt(2.0 * self.config.num_hidden_layers)
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, (nn.Embedding, nn.Conv1d)):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, CometRMSNorm):
            nn.init.ones_(module.weight)
        elif isinstance(module, CometBlock):
            nn.init.constant_(module.mixer_scale, float(self.config.layer_scale_init))
            nn.init.constant_(module.mlp_scale, float(self.config.layer_scale_init))


class CometModel(CometPreTrainedModel):
    """Bare Comet trunk: token embeddings, blocks and the final norm. Returns hidden states."""

    def __init__(self, config: CometConfig):
        super().__init__(config)
        self._build_trunk(config)
        self.post_init()

    def _build_trunk(self, config: CometConfig) -> None:
        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList(CometBlock(config, index) for index in range(config.num_hidden_layers))
        self.norm = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False

        # RoPE tables are built lazily on first use and kept as plain attributes rather than
        # buffers: they are derived constants, and this keeps them out of every loading path
        # (meta-device instantiation, device maps, dtype casts).
        self._rope_tables = None

    def _rope(self, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
        tables = self._rope_tables
        if tables is None or tables[0].device != device:
            config = self.config
            head_dim = config.hidden_size // config.num_attention_heads
            exponent = torch.arange(0, head_dim, 2, dtype=torch.float32, device="cpu") / head_dim
            inv_freq = 1.0 / (config.rope_theta**exponent)
            positions = torch.arange(config.max_position_embeddings, dtype=torch.float32, device="cpu")
            frequencies = torch.outer(positions, inv_freq)
            tables = (frequencies.cos().to(device), frequencies.sin().to(device))
            self._rope_tables = tables
        return tables

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embed

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.embed = value

    def _run_trunk(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
    ):
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Pass exactly one of `input_ids` or `inputs_embeds`.")
        if inputs_embeds is None:
            inputs_embeds = self.embed(input_ids)
        batch, sequence = inputs_embeds.shape[:2]
        device = inputs_embeds.device

        use_cache = self.config.use_cache if use_cache is None else use_cache
        if self.training and past_key_values is None:
            use_cache = False
        cache = past_key_values
        if cache is not None and not hasattr(cache, "update"):
            raise TypeError(
                "`past_key_values` must be a `transformers.Cache` instance (for example `DynamicCache()`); "
                "legacy tuple caches are not supported."
            )
        if cache is None and use_cache:
            cache = DynamicCache()
        past = _cache_length(cache)
        total = past + sequence
        if total > self.config.max_position_embeddings:
            raise ValueError(
                f"Sequence length {total} exceeds max_position_embeddings={self.config.max_position_embeddings}."
            )

        # ---- masks -------------------------------------------------------------
        attn_mask = None  # [B|1, 1, T, past + T] boolean, True = may attend
        token_mask = None  # [B, T] boolean, True = real token (current chunk only)
        is_causal = False
        keep = None
        if attention_mask is not None:
            if attention_mask.dim() != 2:
                raise ValueError("`attention_mask` must have shape [batch, sequence].")
            keep = attention_mask.to(device=device, dtype=torch.bool)
            if keep.shape[1] == sequence and past > 0:
                # Only the new chunk was described; everything cached is treated as real.
                keep = torch.cat([keep.new_ones(batch, past), keep], dim=1)
            if keep.shape[1] != total:
                raise ValueError(
                    f"`attention_mask` has length {attention_mask.shape[1]} but the model sees {total} positions "
                    f"({past} cached + {sequence} new)."
                )
            if bool(keep.all()):
                keep = None

        if keep is not None or (past > 0 and sequence > 1):
            query_positions = past + torch.arange(sequence, device=device)
            key_positions = torch.arange(total, device=device)
            attn_mask = (key_positions[None, :] <= query_positions[:, None])[None, None]
            if keep is not None:
                # Every position may always attend to itself so that padded rows never
                # become fully masked (which would produce NaNs).
                diagonal = (key_positions[None, :] == query_positions[:, None])[None, None]
                attn_mask = (attn_mask & keep[:, None, None, :]) | diagonal
                token_mask = keep[:, past:]
        elif past == 0 and sequence > 1:
            is_causal = True

        # ---- positions ---------------------------------------------------------
        if position_ids is None:
            if keep is not None:
                # Real tokens are numbered 0, 1, 2, ... regardless of how much padding precedes them.
                position_ids = (keep.long().cumsum(-1) - 1).clamp_min(0)[:, past:]
            else:
                position_ids = torch.arange(past, total, device=device)[None, :]
        else:
            position_ids = position_ids.to(device)
            if position_ids.dim() == 1:
                position_ids = position_ids[None, :]
            if position_ids.shape[-1] != sequence:
                position_ids = position_ids[..., -sequence:]
            if int(position_ids.max()) >= self.config.max_position_embeddings:
                raise ValueError("`position_ids` exceed max_position_embeddings.")
        rope_cos, rope_sin = self._rope(device)
        cos = rope_cos[position_ids].unsqueeze(1)
        sin = rope_sin[position_ids].unsqueeze(1)

        # ---- blocks ------------------------------------------------------------
        x = inputs_embeds
        all_hidden_states = () if output_hidden_states else None
        for block in self.blocks:
            if output_hidden_states:
                all_hidden_states += (x,)
            if self.gradient_checkpointing and self.training and cache is None:
                x = torch.utils.checkpoint.checkpoint(
                    block, x, cos, sin, attn_mask, token_mask, is_causal, use_reentrant=False
                )
            else:
                x = block(x, cos, sin, attn_mask=attn_mask, token_mask=token_mask, is_causal=is_causal, cache=cache)
        x = self.norm(x)
        if output_hidden_states:
            all_hidden_states += (x,)
        return x, (cache if use_cache else None), all_hidden_states

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        if output_hidden_states is None:
            output_hidden_states = bool(getattr(self.config, "output_hidden_states", False))
        if return_dict is None:
            return_dict = bool(getattr(self.config, "return_dict", True))
        hidden, cache, all_hidden_states = self._run_trunk(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
        )
        if not return_dict:
            return tuple(item for item in (hidden, cache, all_hidden_states) if item is not None)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden,
            past_key_values=cache,
            hidden_states=all_hidden_states,
        )


class CometForCausalLM(CometModel, GenerationMixin):
    """
    Comet trunk with the language-modelling head.

    The state-dict layout is flat and identical to the training checkpoint
    (``embed.weight``, ``blocks.N.*``, ``norm.weight``). With
    ``config.tie_word_embeddings=True`` there is no separate head parameter: logits are
    computed directly against ``embed.weight``.
    """

    def __init__(self, config: CometConfig):
        CometPreTrainedModel.__init__(self, config)
        self._build_trunk(config)
        if config.tie_word_embeddings:
            self.lm_head = None
        else:
            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_output_embeddings(self):
        # `None` when tied: there is nothing for Transformers to tie or resize separately.
        return self.lm_head

    def set_output_embeddings(self, value) -> None:
        self.lm_head = value

    def _logits(self, hidden: torch.Tensor) -> torch.Tensor:
        weight = self.embed.weight if self.lm_head is None else self.lm_head.weight
        return F.linear(hidden.to(weight.device), weight)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch, sequence)`, *optional*):
            Next-token targets. They are shifted inside the model, so pass ``labels=input_ids``;
            positions set to ``-100`` are ignored.
        """
        if output_hidden_states is None:
            output_hidden_states = bool(getattr(self.config, "output_hidden_states", False))
        if return_dict is None:
            return_dict = bool(getattr(self.config, "return_dict", True))
        hidden, cache, all_hidden_states = self._run_trunk(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
        )
        logits = self._logits(hidden)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1].float()
            shift_labels = labels[:, 1:].to(shift_logits.device)
            total = F.cross_entropy(
                shift_logits.reshape(-1, shift_logits.shape[-1]),
                shift_labels.reshape(-1),
                ignore_index=-100,
                reduction="sum",
            )
            loss = total / shift_labels.ne(-100).sum().clamp_min(1)

        if not return_dict:
            return tuple(item for item in (loss, logits, cache, all_hidden_states) if item is not None)
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=cache,
            hidden_states=all_hidden_states,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values=None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> dict:
        """Feed only the tokens the cache has not seen yet; keep the full-length attention mask."""
        past = _cache_length(past_key_values)
        model_inputs = {}
        if inputs_embeds is not None and past == 0:
            model_inputs["inputs_embeds"] = inputs_embeds
        else:
            seen = past
            if inputs_embeds is not None:
                # Generation started from embeddings: `input_ids` holds only the generated tokens.
                seen = past - inputs_embeds.shape[1]
            unseen = input_ids.shape[1] - seen
            if past > 0:
                input_ids = input_ids[:, -unseen:] if unseen > 0 else input_ids[:, -1:]
            model_inputs["input_ids"] = input_ids.contiguous()
        model_inputs["past_key_values"] = past_key_values
        model_inputs["attention_mask"] = attention_mask
        model_inputs["use_cache"] = kwargs.get("use_cache", True)
        if kwargs.get("output_hidden_states") is not None:
            model_inputs["output_hidden_states"] = kwargs["output_hidden_states"]
        return model_inputs


__all__ = ["CometPreTrainedModel", "CometModel", "CometForCausalLM"]