""" 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"]