import math import torch import torch.nn as nn import torch.nn.functional as F import json from safetensors.torch import save_file, load_file import os from huggingface_hub import PyTorchModelHubMixin from transformers import PretrainedConfig, PreTrainedModel class CosmoFormerConfig(PretrainedConfig): model_type = "cosmoformer" def __init__( self, d_model: int = 256, d_ff: int = 512, dropout: float = 0.1, num_groups: int = 4, num_heads: int = 8, num_layers: int = 6, vocab_size: int = 65400, max_len: int = 2048, **kwargs ): super().__init__(**kwargs) self.d_model = d_model self.d_ff = d_ff self.dropout = dropout self.num_groups = num_groups self.num_heads = num_heads self.num_layers = num_layers self.vocab_size = vocab_size self.max_len = max_len class SinusoidalPositionalEncoding(nn.Module): def __init__(self, d_model: int, max_len: int = 5000): super().__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer('pe', pe.unsqueeze(0)) def forward(self, x: torch.Tensor) -> torch.Tensor: S = x.size(1) return x + self.pe[:, :S, :] class Embedder(nn.Module): def __init__(self, vocab: int, d_model: int): super().__init__() self.emb = nn.Embedding(vocab, d_model) self.d_model = d_model def forward(self, x: torch.Tensor): return self.emb(x)# * math.sqrt(self.d_model) class FFN(nn.Module): def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1): super().__init__() self.linear1 = nn.Linear(d_model, d_ff) self.linear2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) self.activation = nn.GELU() def forward(self, x: torch.Tensor) -> torch.Tensor: out = self.linear1(x) out = self.activation(out) out = self.dropout(out) out = self.linear2(out) return out class GroupedQueryAttention(nn.Module): def __init__(self, d_model, num_heads, num_groups, dropout=0.0): super().__init__() assert d_model % num_heads == 0 assert num_heads % num_groups == 0 self.d_model = d_model self.num_heads = num_heads self.num_groups = num_groups self.head_dim = d_model // num_heads self.heads_per_group = num_heads // num_groups self.q_proj = nn.Linear(d_model, d_model) self.k_proj = nn.Linear(d_model, num_groups * self.head_dim) self.v_proj = nn.Linear(d_model, num_groups * self.head_dim) self.out_proj = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, query, key, value, key_padding_mask=None, is_causal=False): B, S, _ = query.shape q = self.q_proj(query) k = self.k_proj(key) v = self.v_proj(value) q = q.view(B, S, self.num_heads, self.head_dim).transpose(1, 2) k = k.view(B, S, self.num_groups, self.head_dim).transpose(1, 2) v = v.view(B, S, self.num_groups, self.head_dim).transpose(1, 2) k = k.unsqueeze(2).expand(-1, -1, self.heads_per_group, -1, -1).reshape(B, self.num_heads, S, self.head_dim) v = v.unsqueeze(2).expand(-1, -1, self.heads_per_group, -1, -1).reshape(B, self.num_heads, S, self.head_dim) attn_mask = None if is_causal or key_padding_mask is not None: causal_mask = torch.triu(torch.ones(S, S, device=query.device) * float('-inf'), diagonal=1) if key_padding_mask is not None: pad_mask = torch.where(key_padding_mask, float('-inf'), 0.0) pad_mask = pad_mask.unsqueeze(1).unsqueeze(2) # (B,1,1,S) attn_mask = causal_mask + pad_mask else: attn_mask = causal_mask is_causal = False attn_output = F.scaled_dot_product_attention( q, k, v, attn_mask=attn_mask, dropout_p=self.dropout.p if self.training else 0.0, is_causal=is_causal ) attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, self.d_model) return self.out_proj(attn_output) class DecoderLayer(nn.Module): def __init__(self, d_model: int, d_ff: int, dropout: float, num_groups: int, num_heads: int): super().__init__() self.gqa = GroupedQueryAttention(d_model, num_heads, num_groups, dropout) self.ffn = FFN(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, key_padding_mask=None): residual = x x = self.norm1(x) attn_out = self.gqa(query=x, key=x, value=x, key_padding_mask=key_padding_mask, is_causal=True) x = residual + self.dropout(attn_out) residual = x x = self.norm2(x) ff_out = self.ffn(x) x = residual + self.dropout(ff_out) return x class CosmoFormer(PreTrainedModel): config_class = CosmoFormerConfig def __init__(self, config: CosmoFormerConfig): super().__init__(config) self.config = config self.d_model = config.d_model self.vocab_size = config.vocab_size self.num_layers = config.num_layers self.embedder = Embedder(config.vocab_size, config.d_model) self.pe = SinusoidalPositionalEncoding(config.d_model, config.max_len) self.layers = nn.ModuleList([ DecoderLayer(config.d_model, config.d_ff, config.dropout, config.num_groups, config.num_heads) for _ in range(config.num_layers) ]) self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) self._init_weights() self.post_init() def _init_weights(self): for p in self.parameters(): if p.dim() > 1: nn.init.normal_(p, mean=0.0, std=0.02) else: nn.init.zeros_(p) def forward(self, input_ids, attention_mask=None, labels=None): batch, seq_len = input_ids.shape device = input_ids.device x = self.embedder(input_ids) x = self.pe(x) key_padding_mask = None if attention_mask is not None: key_padding_mask = (attention_mask == 0).to(device) for layer in self.layers: x = layer(x, key_padding_mask=key_padding_mask) logits = self.lm_head(x) loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss = F.cross_entropy(shift_logits.view(-1, self.vocab_size), shift_labels.view(-1), ignore_index=-100) return (loss, logits) if loss is not None else logits def generate(self, input_ids, max_new_tokens=50, temperature=1.0, do_sample=False, top_k=None, top_p=None, eos_token_id=None, **kwargs): self.eval() generated = input_ids.clone() for _ in range(max_new_tokens): logits = self.forward(generated) next_logits = logits[:, -1, :] / temperature if top_k is not None and top_k > 0: indices_to_remove = next_logits < torch.topk(next_logits, top_k)[0][:, -1, None] next_logits[indices_to_remove] = float('-inf') if top_p is not None and top_p < 1.0: sorted_logits, sorted_indices = torch.sort(next_logits, descending=True) cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) next_logits[indices_to_remove] = float('-inf') if do_sample: probs = torch.softmax(next_logits, dim=-1) next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) else: next_tokens = torch.argmax(next_logits, dim=-1) generated = torch.cat([generated, next_tokens.unsqueeze(1)], dim=1) if eos_token_id is not None and (next_tokens == eos_token_id).all(): break return generated def num_parameters(self, only_trainable: bool = False) -> int: if only_trainable: return sum(p.numel() for p in self.parameters() if p.requires_grad) return sum(p.numel() for p in self.parameters())