src/ares/model/model.py

"""
AresForCausalLM - Decoder-Only Transformer, from scratch.
No external APIs for core intelligence.
Includes:
- Token embeddings (LOTS)
- RoPE
- MHA+GQA+KV Cache
- RMSNorm
- SwiGLU/GELU
- Unembedding matrix multipliers
- Token outputs / logits
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, List, Tuple
from ..config import AresConfig
from .rmsnorm import RMSNorm
from .transformer_block import AresDecoderLayer
from .attention import KVCache

class AresModel(nn.Module):
    def __init__(self, config: AresConfig):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)  # token embedding matrix
        self.layers = nn.ModuleList([AresDecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False

    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, kv_caches=None):
        hidden_states = self.embed_tokens(input_ids)  # [b,s,hidden]

        presents = [] if use_cache else None
        for i, layer in enumerate(self.layers):
            past_kv = past_key_values[i] if past_key_values is not None else None
            cache = kv_caches[i] if kv_caches is not None else None

            if self.gradient_checkpointing and self.training:
                # Flaw fix: memory blow -> gradient checkpointing
                def custom_forward(h):
                    return layer(h, attention_mask=attention_mask, past_kv=past_kv, use_cache=use_cache, cache=cache)[0]
                hidden_states = torch.utils.checkpoint.checkpoint(custom_forward, hidden_states)
                # present not tracked in checkpoint mode for simplicity
                if use_cache:
                    # dummy
                    pass
            else:
                hidden_states, present = layer(
                    hidden_states, attention_mask=attention_mask, past_kv=past_kv, use_cache=use_cache, cache=cache
                )
                if use_cache:
                    presents.append(present)

        hidden_states = self.norm(hidden_states)
        return hidden_states, presents

class AresForCausalLM(nn.Module):
    def __init__(self, config: AresConfig):
        super().__init__()
        self.config = config
        self.model = AresModel(config)
        # Unembedding matrix multiplier
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)  # output projection

        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
        use_cache: bool = False,
        kv_caches: Optional[List[KVCache]] = None,
    ):
        hidden_states, presents = self.model(
            input_ids, attention_mask=attention_mask, past_key_values=past_key_values, use_cache=use_cache, kv_caches=kv_caches
        )
        logits = self.lm_head(hidden_states)  # token sections -> token outputs [b,s,vocab]

        loss = None
        if labels is not None:
            # Shift for causal LM
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))

        return {"logits": logits, "loss": loss, "past_key_values": presents, "hidden_states": hidden_states}

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 0.8,
        top_p: float = 0.9,
        top_k: int = 50,
        do_sample: bool = True,
        eos_token_id: Optional[int] = None,
    ):
        self.eval()
        bsz, seq_len = input_ids.shape
        device = input_ids.device
        # Initialize KV caches for fast inference
        kv_caches = [
            KVCache(bsz, self.config.num_key_value_heads, self.config.max_position_embeddings, self.config.head_dim, device, self.model.embed_tokens.weight.dtype)
            for _ in range(self.config.num_hidden_layers)
        ]

        # Prefill
        hidden, presents = self.model(input_ids[:, :-1], use_cache=False) if seq_len > 1 else (None, None)
        # Actually for cache prefill, we need to run first tokens through cache
        # Simplify: run all but last token to fill cache
        if seq_len > 1:
            # Fill cache iteratively - for simplicity use past_key_values not KVCache class here, but we will use cache
            # Prefill using kv_caches
            # First token
            cur_ids = input_ids[:, :1]
            for i in range(1, seq_len):
                _ = self.model(cur_ids[:, -1:], past_key_values=None, kv_caches=kv_caches if i>1 else None, use_cache=False)
                # Actually our model forward uses kv_caches list - let's call layer-wise? easier: just loop model forward for each token with kv_caches
                # We need to handle cache update inside attention, so we call model forward each step
                # Reset approach: use past_key_values accumulation instead of KVCache class for generate simplicity
                pass
            # For simplicity, abandon KVCache class in generate and use past_key_values
            past = None
            # Re-prefill correctly with past_key_values
            if seq_len > 1:
                out = self.forward(input_ids[:, :-1], use_cache=True)
                past = out["past_key_values"]
                logits = out["logits"][:, -1:, :]
            else:
                out = self.forward(input_ids, use_cache=True)
                past = out["past_key_values"]
                logits = out["logits"][:, -1:, :]
        else:
            out = self.forward(input_ids, use_cache=True)
            past = out["past_key_values"]
            logits = out["logits"][:, -1:, :]

        generated = input_ids

        for _ in range(max_new_tokens):
            # Last token logits
            last_logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)

            if do_sample:
                if top_k > 0:
                    top_k_vals, _ = torch.topk(last_logits, top_k)
                    min_top_k = top_k_vals[:, -1].unsqueeze(-1)
                    last_logits = torch.where(last_logits < min_top_k, torch.full_like(last_logits, float('-inf')), last_logits)
                if top_p < 1.0:
                    sorted_logits, sorted_indices = torch.sort(last_logits, descending=True)
                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                    sorted_mask = cumulative_probs > top_p
                    sorted_mask[:, 1:] = sorted_mask[:, :-1].clone()
                    sorted_mask[:, 0] = 0
                    indices_to_remove = sorted_mask.scatter(1, sorted_indices, sorted_mask)
                    last_logits = last_logits.masked_fill(indices_to_remove, float('-inf'))
                probs = F.softmax(last_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(last_logits, dim=-1, keepdim=True)

            generated = torch.cat([generated, next_token], dim=1)

            if eos_token_id is not None and (next_token == eos_token_id).all():
                break

            # next forward with cache
            out = self.forward(next_token, past_key_values=past, use_cache=True)
            past = out["past_key_values"]
            logits = out["logits"]

            if generated.shape[1] >= self.config.max_position_embeddings:
                break

        return generated

    def count_parameters(self):
        return sum(p.numel() for p in self.parameters())