Delete transformer.py

transformer.py

@@ -1,90 +0,0 @@
-import torch
-import torch.nn as nn
-
-from .config import IvmeConfig
-from .rmsnorm import RMSNorm
-from .rope import precompute_rope_freqs
-from .attention import CausalSelfAttention
-from .feedforward import SwiGLU
-
-
-class TransformerBlock(nn.Module):
-    """One dense transformer layer (Section 3): pre-norm attention + pre-norm SwiGLU,
-    with residual connections around each. Identical shape repeated n_layers times --
-    no loops, no weight sharing (Section 3.1, distinguishing this from the shelved
-    Ivmetron design).
-    """
-
-    def __init__(self, cfg: IvmeConfig):
-        super().__init__()
-        self.attn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
-        self.attn = CausalSelfAttention(cfg.hidden_dim, cfg.n_heads, cfg.dropout)
-        self.ffn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
-        self.ffn = SwiGLU(cfg.hidden_dim, cfg.ffn_mult)
-
-    def forward(self, x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor:
-        x = x + self.attn(self.attn_norm(x), rope_freqs)
-        x = x + self.ffn(self.ffn_norm(x))
-        return x
-
-
-class IvmeConversateV2(nn.Module):
-    """Ivme-Conversate-v2 (Dense) -- the full model described in Section 4.
-
-    ~20M parameters, 10 layers, hidden_dim 384, 6 heads, RoPE, SwiGLU, RMSNorm,
-    tied embeddings, 16k vocab, 1024 context. See config.py for the exact spec.
-    """
-
-    def __init__(self, cfg: IvmeConfig):
-        super().__init__()
-        self.cfg = cfg
-
-        self.tok_embed = nn.Embedding(cfg.vocab_size, cfg.hidden_dim)
-        self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])
-        self.final_norm = RMSNorm(cfg.hidden_dim, eps=cfg.norm_eps)
-
-        # Section 4.8: tied embeddings -- output head reuses the input embedding
-        # table instead of learning a separate one.
-        self.lm_head = nn.Linear(cfg.hidden_dim, cfg.vocab_size, bias=False)
-        if cfg.tie_embeddings:
-            self.lm_head.weight = self.tok_embed.weight
-
-        rope_freqs = precompute_rope_freqs(cfg.head_dim, cfg.context_len, cfg.rope_theta)
-        self.register_buffer("rope_freqs", rope_freqs, persistent=False)
-
-        self.apply(self._init_weights)
-
-    def _init_weights(self, module: nn.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: torch.Tensor, targets: torch.Tensor | None = None):
-        B, T = idx.shape
-        assert T <= self.cfg.context_len, (
-            f"sequence length {T} exceeds context_len {self.cfg.context_len}"
-        )
-
-        x = self.tok_embed(idx)
-        for block in self.blocks:
-            x = block(x, self.rope_freqs)
-        x = self.final_norm(x)
-        logits = self.lm_head(x)
-
-        loss = None
-        if targets is not None:
-            loss = nn.functional.cross_entropy(
-                logits.view(-1, logits.size(-1)),
-                targets.view(-1),
-                ignore_index=-1,
-            )
-        return logits, loss
-
-    def num_params(self, non_embedding: bool = False) -> int:
-        n = sum(p.numel() for p in self.parameters())
-        if non_embedding:
-            n -= self.tok_embed.weight.numel()
-        return n