Mini-GPT: End-to-End Language Model

hard · gpt, transformers, language-models

Mini-GPT: End-to-End Language Model

Build a GPT-style model with multi-head attention, train it with cross-entropy, and generate text autoregressively.

Components to implement

1) `Embedding(Module)`

Token and positional embedding lookup tables.

2) `LayerNorm(Module)`

Normalize a single vector of length D.

3) `FeedForward(Module)`

Two linear layers with ReLU activation (D -> 4D -> D).

4) `MultiHeadAttention(Module)`

Split into heads, run attention per head, concatenate, and project.

5) `TransformerBlock(Module)`

Pre-norm block with residuals:

y = x + MHA(LN(x))
z = y + FFN(LN(y))

6) `MiniGPT(Module)`

Complete model:

token embedding
position embedding
num_layers transformer blocks
final LayerNorm
Linear projection to vocab size

7) Helpers

softmax
cross_entropy_loss
generate

Requirements

Input token_ids length T must be <= max_seq_len.
Output logits shape: (T, vocab_size).
cross_entropy_loss averages over positions.
generate crops context to max_seq_len and samples next token.
embed_dim must be divisible by num_heads.

from __future__ import annotations
import math
import random
from typing import List, Iterator, Tuple, Callable, Set

class Value:
    """Scalar Value with autograd support."""

def __init__(self, data: float, _children: Tuple["Value", ...] = (), _op: str = "") -> None:
        self.data = float(data)
        self.grad = 0.0
        self._backward: Callable[[], None] = lambda: None
        self._prev: Set["Value"] = set(_children)
        self._op = _op

def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")
        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")
        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

def __pow__(self, n):
        out = Value(self.data ** n, (self,), f"**{n}")
        def _backward():
            self.grad += n * (self.data ** (n - 1)) * out.grad
        out._backward = _backward
        return out

def __neg__(self):
        return self * -1

def __sub__(self, other):
        return self + (-other)

def __truediv__(self, other):
        return self * (other ** -1)

def __radd__(self, other):
        return self + other

def __rmul__(self, other):
        return self * other

def exp(self):
        out = Value(math.exp(self.data), (self,), "exp")
        def _backward():
            self.grad += out.data * out.grad
        out._backward = _backward
        return out

def log(self):
        out = Value(math.log(max(self.data, 1e-10)), (self,), "log")
        def _backward():
            self.grad += (1.0 / max(self.data, 1e-10)) * out.grad
        out._backward = _backward
        return out

def relu(self):
        out = Value(max(0.0, self.data), (self,), "relu")
        def _backward():
            self.grad += (1.0 if self.data > 0 else 0.0) * out.grad
        out._backward = _backward
        return out

def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

class Module:
    """Module base class."""

def __init__(self) -> None:
        object.__setattr__(self, "_parameters", {})
        object.__setattr__(self, "_modules", {})
        object.__setattr__(self, "training", True)

def __setattr__(self, name: str, value) -> None:
        if isinstance(value, Value):
            self._parameters[name] = value
        elif isinstance(value, Module):
            self._modules[name] = value
        elif isinstance(value, (list, tuple)):
            self._register_list(name, value)
        object.__setattr__(self, name, value)

def _register_list(self, name: str, values) -> None:
        for i, v in enumerate(values):
            if isinstance(v, Value):
                self._parameters[f"{name}.{i}"] = v
            elif isinstance(v, Module):
                self._modules[f"{name}.{i}"] = v
            elif isinstance(v, (list, tuple)):
                self._register_list(f"{name}.{i}", v)

def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

def forward(self, *args, **kwargs):
        raise NotImplementedError()

def parameters(self) -> List[Value]:
        params = list(self._parameters.values())
        for module in self._modules.values():
            params.extend(module.parameters())
        return params

class Linear(Module):
    """Linear layer."""

def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
        super().__init__()
        k = 1.0 / (in_features ** 0.5)
        self.weight = [
            [Value(random.uniform(-k, k)) for _ in range(in_features)]
            for _ in range(out_features)
        ]
        self.bias = [Value(0.0) for _ in range(out_features)] if bias else None
        self.out_features = out_features

def forward(self, x: List[Value]) -> List[Value]:
        out = []
        for i in range(self.out_features):
            val = sum(w * xi for w, xi in zip(self.weight[i], x))
            if self.bias is not None:
                val = val + self.bias[i]
            out.append(val)
        return out

def softmax(scores: List[Value]) -> List[Value]:
    """TODO: Implement numerically stable softmax."""
    raise NotImplementedError()

def attention_scores(Q: List[List[Value]], K: List[List[Value]], d_k: int) -> List[List[Value]]:
    seq_len = len(Q)
    scale = d_k ** 0.5
    scores = []
    for i in range(seq_len):
        row = []
        for j in range(seq_len):
            dot = sum(q * k for q, k in zip(Q[i], K[j]))
            row.append(dot / scale)
        scores.append(row)
    return scores

def causal_mask(seq_len: int) -> List[List[float]]:
    mask = []
    for i in range(seq_len):
        row = []
        for j in range(seq_len):
            row.append(0.0 if j <= i else -1e9)
        mask.append(row)
    return mask

def apply_attention(weights: List[List[Value]], V: List[List[Value]]) -> List[List[Value]]:
    seq_len = len(weights)
    d_v = len(V[0])
    output = []
    for i in range(seq_len):
        attended = []
        for k in range(d_v):
            val = Value(0.0)
            for j in range(seq_len):
                val = val + weights[i][j] * V[j][k]
            attended.append(val)
        output.append(attended)
    return output

class Embedding(Module):
    """TODO: Implement embedding lookup."""

def __init__(self, num_embeddings: int, embedding_dim: int) -> None:
        raise NotImplementedError()

def forward(self, indices: List[int]) -> List[List[Value]]:
        raise NotImplementedError()

class LayerNorm(Module):
    """TODO: Implement layer normalization."""

def __init__(self, dim: int, eps: float = 1e-5) -> None:
        raise NotImplementedError()

def forward(self, x: List[Value]) -> List[Value]:
        raise NotImplementedError()

class FeedForward(Module):
    """TODO: Implement feed-forward network."""

def __init__(self, embed_dim: int, hidden_dim: int = None) -> None:
        raise NotImplementedError()

def forward(self, x: List[Value]) -> List[Value]:
        raise NotImplementedError()

class MultiHeadAttention(Module):
    """TODO: Implement multi-head attention."""

def __init__(self, embed_dim: int, num_heads: int) -> None:
        raise NotImplementedError()

def forward(self, x: List[List[Value]], causal: bool = True) -> List[List[Value]]:
        raise NotImplementedError()

class TransformerBlock(Module):
    """TODO: Implement transformer block."""

def __init__(self, embed_dim: int, num_heads: int) -> None:
        raise NotImplementedError()

def forward(self, x: List[List[Value]]) -> List[List[Value]]:
        raise NotImplementedError()

class MiniGPT(Module):
    """TODO: Implement complete GPT model."""

def __init__(
        self,
        vocab_size: int,
        embed_dim: int,
        num_heads: int,
        num_layers: int,
        max_seq_len: int,
    ) -> None:
        raise NotImplementedError()

def forward(self, token_ids: List[int]) -> List[List[Value]]:
        raise NotImplementedError()

def cross_entropy_loss(logits: List[List[Value]], targets: List[int]) -> Value:
    """TODO: Mean cross-entropy over positions."""
    raise NotImplementedError()

def generate(
    model: MiniGPT,
    start_ids: List[int],
    max_new_tokens: int,
    temperature: float = 1.0,
) -> List[int]:
    """TODO: Autoregressive sampling."""
    raise NotImplementedError()