Update modeling_tinygpt.py

modeling_tinygpt.py

@@ -1,63 +1,86 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from torch.nn import CrossEntropyLoss
 from transformers.modeling_outputs import CausalLMOutput
 from transformers.modeling_utils import PreTrainedModel
 from configuration_tinygpt import TinyGPTConfig # Changed from relative to absolute import
 
+# -------------------------
+# TinyGPTConfig (Required)
+# -------------------------
+class TinyGPTConfig:
+    model_type = "tinygpt"
+
+    def __init__(self,
+                 vocab_size=30522,
+                 d_model=256,
+                 n_heads=4,
+                 n_layers=4,
+                 d_ff=1024,
+                 max_seq_len=256,
+                 **kwargs):
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.d_ff = d_ff
+        self.max_seq_len = max_seq_len
+
+        # store additional HF keys
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+
+# -------------------------
+# Your Original TinyGPT Core
+# -------------------------
+class TinyGPT(nn.Module):
+    def __init__(self, vocab_size=30522, d_model=256, n_heads=4,
+                 n_layers=4, d_ff=1024, max_seq_len=256):
+        x = self.ln_f(x)
+        return self.head(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff):
+        super().__init__()
+        ff_out = self.ff(x)
+        x = self.ln2(x + ff_out)
+        return x
+
+
+# -------------------------
+# HF Wrapper: TinyGPTForCausalLM
+# -------------------------
 class TinyGPTForCausalLM(PreTrainedModel):
     config_class = TinyGPTConfig
 
     def __init__(self, config):
         super().__init__(config)
 
-        self.embed = nn.Embedding(config.vocab_size, config.d_model)
-        self.pos_embed = nn.Embedding(config.max_seq_len, config.d_model)
-
-        self.blocks = nn.ModuleList([
-            nn.TransformerEncoderLayer(
-                d_model=config.d_model,
-                nhead=config.n_heads,
-                dim_feedforward=config.d_ff,
-                batch_first=True
-            )
-            for _ in range(config.n_layers)
-        ])
-
-        self.norm = nn.LayerNorm(config.d_model)
-        self.lm_head = nn.Linear(config.d_model, config.vocab_size)
+        self.model = TinyGPT(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            n_heads=config.n_heads,
+            n_layers=config.n_layers,
+            d_ff=config.d_ff,
+            max_seq_len=config.max_seq_len
+        )
 
         self.post_init()
 
     def forward(self, input_ids, labels=None):
-        B, T = input_ids.shape
-        positions = torch.arange(T, device=input_ids.device).unsqueeze(0)
-
-        x = self.embed(input_ids) + self.pos_embed(positions)
-
-        for blk in self.blocks:
-            x = blk(x)
-
-        x = self.norm(x)
-        logits = self.lm_head(x)
+        logits = self.model(input_ids)
 
         loss = None
         if labels is not None:
-            shift_logits = logits[:, :-1, :].contiguous()
-            shift_labels = labels[:, 1:].contiguous()
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)),
-                            shift_labels.view(-1))
+            loss = nn.CrossEntropyLoss()(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1)
+            )
 
         return CausalLMOutput(
-            loss=loss,
-            logits=logits
-        )
-
-    @torch.no_grad()
-    def generate(self, input_ids, max_new_tokens=50):
-        for _ in range(max_new_tokens):
-            logits = self.forward(input_ids).logits
-            next_token = torch.argmax(logits[:, -1, :], dim=-1)
-            input_ids = torch.cat([input_ids, next_token[:, None]], dim=1)
-        return input_ids
+            logits=logits,
+            loss=loss
+        )