Update modeling_tinygpt.py

modeling_tinygpt.py

@@ -1,119 +1,63 @@
-
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-from transformers import PreTrainedModel
+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):
-        super().__init__()
-        self.tok_emb = nn.Embedding(vocab_size, d_model)
-        self.pos_emb = nn.Embedding(max_seq_len, d_model)
-
-        self.layers = nn.ModuleList([
-            TransformerBlock(d_model, n_heads, d_ff)
-            for _ in range(n_layers)
-        ])
-
-        self.ln_f = nn.LayerNorm(d_model)
-        self.head = nn.Linear(d_model, vocab_size, bias=False)
-        self.max_seq_len = max_seq_len
-
-    def forward(self, input_ids):
-        b, t = input_ids.size()
-        pos = torch.arange(0, t, device=input_ids.device).unsqueeze(0)
-        x = self.tok_emb(input_ids) + self.pos_emb(pos)
-
-        for layer in self.layers:
-            x = layer(x)
-
-        x = self.ln_f(x)
-        return self.head(x)
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, d_model, n_heads, d_ff):
-        super().__init__()
-        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
-        self.ln1 = nn.LayerNorm(d_model)
-        self.ff = nn.Sequential(
-            nn.Linear(d_model, d_ff),
-            nn.GELU(),
-            nn.Linear(d_ff, d_model)
-        )
-        self.ln2 = nn.LayerNorm(d_model)
-
-    def forward(self, x):
-        attn_out, _ = self.attn(x, x, x)
-        x = self.ln1(x + attn_out)
-        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.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.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.post_init()
 
     def forward(self, input_ids, labels=None):
-        logits = self.model(input_ids)
+        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)
 
         loss = None
         if labels is not None:
-            loss = nn.CrossEntropyLoss()(
-                logits.view(-1, logits.size(-1)),
-                labels.view(-1)
-            )
+            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))
 
         return CausalLMOutput(
-            logits=logits,
-            loss=loss
+            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