Make GPT inherit from PreTrainedModel for HF compatibility

README.md

@@ -13,7 +13,6 @@ library_name: transformers
 pipeline_tag: text-generation
 ---
 
-# 🍼 BabyLangModel
 
 # 🍼 BabyLangModel
 

model.py

@@ -2,6 +2,30 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
+from transformers import PreTrainedModel, PretrainedConfig
+
+class GPTConfig(PretrainedConfig):
+    model_type = "gpt"
+
+    def __init__(
+        self,
+        vocab_size=50257,
+        block_size=128,
+        n_layer=6,
+        n_head=6,
+        n_embd=384,
+        dropout=0.0,
+        bias=True,
+        **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.dropout = dropout
+        self.bias = bias
 
 class LayerNorm(nn.Module):
     def __init__(self, ndim, bias):
@@ -70,10 +94,11 @@ class Block(nn.Module):
         x = x + self.mlp(self.ln2(x))
         return x
 
-class GPT(nn.Module):
+class GPT(PreTrainedModel):
+    config_class = GPTConfig
+
     def __init__(self, config):
-        super().__init__()
-        self.config = config
+        super().__init__(config)
         self.transformer = nn.ModuleDict(dict(
             wte=nn.Embedding(config.vocab_size, config.n_embd),
             wpe=nn.Embedding(config.block_size, config.n_embd),
@@ -97,37 +122,37 @@ class GPT(nn.Module):
         elif isinstance(module, nn.Embedding):
             nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
-    def forward(self, idx, targets=None):
-        device = idx.device
-        b, t = idx.size()
+    def forward(self, input_ids, labels=None):
+        device = input_ids.device
+        b, t = input_ids.size()
         assert t <= self.config.block_size
         pos = torch.arange(0, t, dtype=torch.long, device=device)
 
-        tok_emb = self.transformer.wte(idx)
+        tok_emb = self.transformer.wte(input_ids)
         pos_emb = self.transformer.wpe(pos)
         x = self.transformer.drop(tok_emb + pos_emb)
         for block in self.transformer.h:
             x = block(x)
         x = self.transformer.ln_f(x)
 
-        if targets is not None:
+        if labels is not None:
             logits = self.lm_head(x)
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
-            return logits, loss
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-1)
+            return {'logits': logits, 'loss': loss}
         else:
             logits = self.lm_head(x[:, [-1], :])
-            return logits, None
+            return {'logits': logits}
 
     @torch.no_grad()
-    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+    def generate(self, input_ids, max_new_tokens, temperature=1.0, top_k=None):
         for _ in range(max_new_tokens):
-            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
-            logits, _ = self(idx_cond)
-            logits = logits[:, -1, :] / temperature
+            idx_cond = input_ids if input_ids.size(1) <= self.config.block_size else input_ids[:, -self.config.block_size:]
+            out = self(idx_cond)
+            logits = out['logits'][:, -1, :] / temperature
             if top_k is not None:
                 v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                 logits[logits < v[:, [-1]]] = -float('Inf')
             probs = F.softmax(logits, dim=-1)
             idx_next = torch.multinomial(probs, num_samples=1)
-            idx = torch.cat((idx, idx_next), dim=1)
-        return idx
+            input_ids = torch.cat((input_ids, idx_next), dim=1)
+        return input_ids