Publish Tiny GPT model

README.md

@@ -0,0 +1,62 @@
+---
+language:
+  - en
+license: mit
+library_name: transformers
+pipeline_tag: text-generation
+datasets:
+  - roneneldan/TinyStories
+tags:
+  - custom_code
+  - educational
+---
+
+# Tiny GPT
+
+Tiny GPT is an educational decoder-only Transformer trained from scratch on
+the [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)
+dataset. The implementation is intentionally small and readable.
+
+## Model details
+
+- Architecture: decoder-only causal language model
+- Context length: 512 tokens
+- Vocabulary size: 10,000
+- Hidden size: 256
+- Transformer layers: 6
+- Attention heads: 8
+
+Source code: https://github.com/alainbrown/tiny-gpt
+
+## Usage
+
+This repository contains custom Transformers code. Review it before enabling
+`trust_remote_code`.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo_id = "alainbrown/tiny-gpt"
+tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True)
+
+inputs = tokenizer("Once upon a time", return_tensors="pt")
+logits = model(**inputs).logits
+```
+
+## Intended use
+
+This model is intended for education and experimentation. It is not intended
+for production, factual question answering, or safety-critical applications.
+
+## Limitations
+
+The model is small, trained on synthetic children's stories, and has not been
+comprehensively evaluated. It may produce incoherent, repetitive, incorrect,
+or inappropriate text. English is the only supported language.
+
+## Training
+
+The training pipeline is available in the linked GitHub repository. This model
+repository excludes optimizer and progress state and contains inference files
+only.

config.json

@@ -0,0 +1,25 @@
+{
+  "architectures": [
+    "TinyGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_tiny_gpt.TinyGPTConfig",
+    "AutoModelForCausalLM": "modeling_tiny_gpt.TinyGPTForCausalLM"
+  },
+  "context_size": 512,
+  "d_model": 256,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": 0,
+  "hidden_size": 256,
+  "max_position_embeddings": 512,
+  "model_type": "tiny_gpt",
+  "n_heads": 8,
+  "n_layers": 6,
+  "num_attention_heads": 8,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.12.1",
+  "vocab_size": 10000
+}

configuration_tiny_gpt.py

@@ -0,0 +1,33 @@
+from transformers import PretrainedConfig
+
+
+class TinyGPTConfig(PretrainedConfig):
+    model_type = "tiny_gpt"
+
+    def __init__(
+        self,
+        context_size=32,
+        vocab_size=1024,
+        d_model=64,
+        n_layers=4,
+        n_heads=4,
+        dropout=0.1,
+        tie_word_embeddings=True,
+        use_cache=False,
+        **kwargs,
+    ):
+        self.context_size = context_size
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.hidden_size = d_model
+        self.num_hidden_layers = n_layers
+        self.num_attention_heads = n_heads
+        self.max_position_embeddings = context_size
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            use_cache=use_cache,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "5.12.1"
+}

model.py

@@ -0,0 +1,166 @@
+import math
+import torch
+from torch import nn
+
+"""
+token embeddings
+learned positional embeddings
+causal self-attention
+feed-forward network
+custom LayerNorm
+residual connections
+stacked transformer blocks
+GELU
+dropout
+Pre-LayerNorm
+tied token embedding/output projection weights
+multi-head attention
+
+"""
+class Model(nn.Module):
+    def __init__(self, context_size, vocab_size, d_model, n_layers, n_heads, dropout=0.1):
+        super().__init__()
+
+        self.context_size = context_size
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.dropout_p = dropout
+
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        self.position_embedding = nn.Embedding(context_size, d_model)
+        self.transformer_blocks = nn.ModuleList(
+            [TransformerBlock(d_model, n_heads, dropout) for _ in range(n_layers)]
+        )
+        self.linear = nn.Linear(d_model, vocab_size, bias=False)
+        self.dropout = nn.Dropout(dropout)
+        self.final_layer_norm = LayerNorm(d_model)
+
+    def forward(self, x):
+        B, T = x.shape
+
+        assert T <= self.context_size, "Input sequence is longer than context_size"
+
+        positions = torch.arange(T, device=x.device)
+
+        position = self.position_embedding(positions)
+        token = self.token_embedding(x)
+
+        x = token + position
+        x = self.dropout(x)
+
+        for block in self.transformer_blocks:
+            x = block(x)
+
+        x = self.final_layer_norm(x)
+        logits = self.linear(x)
+
+        return logits
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.ff1 = nn.Linear(d_model, 4 * d_model)
+        self.ff2 = nn.Linear(d_model * 4, d_model)
+
+    def forward(self, x):
+        x = self.ff1(x)
+        x = nn.functional.gelu(x)
+        x = self.ff2(x)
+        return x
+
+class LayerNorm(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(d_model))
+        self.beta = nn.Parameter(torch.zeros(d_model))
+
+    def forward(self, x):
+        mean = x.mean(dim=-1, keepdim=True)
+        diff = (x - mean)
+        variance = (diff * diff).mean(dim=-1, keepdim=True)
+        normalized = diff / torch.sqrt(variance + 1e-6)
+        return self.gamma * normalized + self.beta
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model, n_heads, dropout=0.1):
+        super().__init__()
+
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.scale = math.sqrt(self.head_dim)
+
+        self.query = nn.Linear(d_model, d_model)
+        self.key = nn.Linear(d_model, d_model)
+        self.value = nn.Linear(d_model, d_model)
+
+        self.head_proj = nn.Linear(d_model, d_model)
+
+    def forward(self, x):
+        query = self.split_heads(self.query(x))
+        key = self.split_heads(self.key(x))
+        value = self.split_heads(self.value(x))
+
+        scores = torch.matmul(query, key.transpose(-2, -1))
+        scores = scores / self.scale
+
+        context_size = query.shape[2]
+
+        mask = torch.tril(
+            torch.ones(context_size, context_size, device=query.device)
+        )
+        mask = mask.view(1, 1, context_size, context_size)
+
+        scores = scores.masked_fill(mask == 0, float("-inf"))
+
+        weights = torch.nn.functional.softmax(scores, dim=-1)
+
+        attended = torch.matmul(weights, value)
+
+        attended = self.combine_heads(attended)
+
+        attended = self.head_proj(attended)
+
+        return attended
+
+    def split_heads(self, x):
+        batch_size, seq_len, d_model = x.shape
+
+        x = x.reshape(batch_size, seq_len, self.n_heads, self.head_dim)
+        x = x.transpose(1, 2)
+
+        return x
+
+    def combine_heads(self, x):
+        batch_size, n_heads, seq_len, head_dim = x.shape
+
+        x = x.transpose(1, 2)
+        x = x.contiguous().view(batch_size, seq_len, n_heads * head_dim)
+
+        return x
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_heads, dropout=0.1):
+        super().__init__()
+
+        self.feed_forward = FeedForward(d_model)
+        self.layer_norm1 = LayerNorm(d_model)
+        self.layer_norm2 = LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.multi_head_attention = MultiHeadAttention(d_model, n_heads, dropout)
+
+    def forward(self, x):
+        attention = self.multi_head_attention(self.layer_norm1(x))
+        attention = self.dropout(attention)
+
+        x = x + attention
+
+        feed_forward = self.feed_forward(self.layer_norm2(x))
+        feed_forward = self.dropout(feed_forward)
+
+        x = x + feed_forward
+
+        return x

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e9f8e95c7fa765cda29765de8ccd1c89d0f08545d42c72943b9c546665db657
+size 39973360

modeling_tiny_gpt.py

@@ -0,0 +1,52 @@
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutput
+
+from .configuration_tiny_gpt import TinyGPTConfig
+from .model import Model
+
+
+class TinyGPTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = TinyGPTConfig
+    main_input_name = "input_ids"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.core_model = Model(
+            context_size=config.context_size,
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            n_layers=config.n_layers,
+            n_heads=config.n_heads,
+            dropout=config.dropout,
+        )
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.core_model.token_embedding
+
+    def set_input_embeddings(self, value):
+        self.core_model.token_embedding = value
+
+    def get_output_embeddings(self):
+        return self.core_model.linear
+
+    def set_output_embeddings(self, new_embeddings):
+        self.core_model.linear = new_embeddings
+
+    def forward(self, input_ids=None, labels=None, **kwargs):
+        if input_ids is None:
+            raise ValueError("input_ids must be provided")
+
+        logits = self.core_model(input_ids)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+
+        return CausalLMOutput(loss=loss, logits=logits)

tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "backend": "tokenizers",
+  "eos_token": "<EOS>",
+  "model_max_length": 512,
+  "pad_token": "<EOS>",
+  "tokenizer_class": "TokenizersBackend"
+}