Update modeling_tinybuddy.py

modeling_tinybuddy.py

@@ -6,10 +6,37 @@ import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PreTrainedModel
-from configuration_tinybuddy import GPTConfig
+from transformers import PreTrainedModel, PretrainedConfig
 
 
+# ========== CONFIG CLASS (embedded to avoid import issues) ==========
+class GPTConfig(PretrainedConfig):
+    model_type = "tinybuddy"
+    
+    def __init__(
+        self,
+        vocab_size: int = 50000,
+        block_size: int = 512,
+        n_layer: int = 6,
+        n_head: int = 8,
+        n_embd: int = 256,
+        mlp_ratio: int = 4,
+        dropout: float = 0.0,
+        tie_weights: bool = False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.mlp_ratio = mlp_ratio
+        self.dropout = dropout
+        self.tie_weights = tie_weights
+
+
+# ========== MODEL ARCHITECTURE ==========
 class CausalSelfAttention(nn.Module):
     def __init__(self, cfg: GPTConfig):
         super().__init__()
@@ -20,7 +47,6 @@ class CausalSelfAttention(nn.Module):
         self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=True)
         self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=True)
         self.drop = nn.Dropout(cfg.dropout)
-        # causal mask
         mask = torch.tril(torch.ones(cfg.block_size, cfg.block_size)).bool()
         self.register_buffer("mask", mask, persistent=False)
 
@@ -31,7 +57,6 @@ class CausalSelfAttention(nn.Module):
         q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
         k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
         v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
-        # use PyTorch's fused SDPA (faster on CPU than manual)
         y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                            dropout_p=self.drop.p if self.training else 0.0)
         y = y.transpose(1, 2).contiguous().view(B, T, C)
@@ -82,23 +107,15 @@ class TinyGPT(PreTrainedModel):
 
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02)
+            module.weight.data.normal_(mean=0.0, std=0.02)
             if module.bias is not None:
                 module.bias.data.zero_()
         elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, 'initializer_range') else 0.02)
-
-    def num_params(self, non_embedding=False):
-        n = sum(p.numel() for p in self.parameters())
-        if non_embedding:
-            n -= self.tok_emb.weight.numel() + self.pos_emb.weight.numel()
-            if not self.config.tie_weights:
-                n -= self.lm_head.weight.numel()
-        return n
+            module.weight.data.normal_(mean=0.0, std=0.02)
 
     def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
         B, T = input_ids.shape
-        assert T <= self.config.block_size, f"sequence length {T} > block_size {self.config.block_size}"
+        assert T <= self.config.block_size
         pos = torch.arange(T, device=input_ids.device)
         x = self.tok_emb(input_ids) + self.pos_emb(pos)[None, :, :]
         x = self.drop(x)
@@ -128,10 +145,7 @@ class TinyGPT(PreTrainedModel):
 
 
 if __name__ == "__main__":
-    from configuration_tinybuddy import GPTConfig
     cfg = GPTConfig()
     m = TinyGPT(cfg)
-    total = m.num_params()
-    nonemb = m.num_params(non_embedding=True)
-    print(f"Total params       : {total:,}  (~{total/1e6:.2f}M)")
-    print(f"Non-embedding params: {nonemb:,}  (~{nonemb/1e6:.2f}M)")
+    total = sum(p.numel() for p in m.parameters())
+    print(f"Total params: {total:,} (~{total/1e6:.2f}M)")