model.py

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())