"""
Abyss-1 the full de coder-only transformer.

    tokens -> embedding -> [block] x n_layer -> final RMSNorm -> LM head

forward() returns logits, and (when targets are given) the cross-entropy
loss for next-token prediction.
"""

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

from .norm import AbyssRMSNorm
from .block import AbyssBlock

8934
class TheAbyssModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.token_emb = nn.Embedding(config.vocab_size, config.n_embd)
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([AbyssBlock(config) for _ in range(config.n_layer)])
        self.final_norm = AbyssRMSNorm(config.n_embd, config.norm_eps)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=True)

        if config.tie_embeddings:
            self.lm_head.weight = self.token_emb.weight

        self.apply(self._init_weights)
        # GPT-2 trick: shrink the projections that feed residual adds, so the
        # residual stream doesn't blow up as depth grows
        for name, p in self.named_parameters():
            if name.endswith("proj.weight") or name.endswith("down.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / (2 * config.n_layer) ** 0.5)

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

    def forward(self, idx, targets=None):
        B, T = idx.shape
        assert T <= self.config.block_size, f"sequence {T} > block_size {self.config.block_size}"

        x = self.drop(self.token_emb(idx))
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,
            )
        return logits, loss
    
    def num_params(self, non_embedding=False):
        n = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n -= self.token_emb.weight.numel()
            if not self.config.tie_embeddings:
                n -= self.lm_head.weight .numel()
        return n