|
|
--- |
|
|
license: mit |
|
|
datasets: |
|
|
- IsmaelMousa/movies |
|
|
tags: |
|
|
- movie |
|
|
- short_stories |
|
|
- llm |
|
|
- slm |
|
|
--- |
|
|
# 🚀 Small Language Model (SLM) from Scratch — Explained |
|
|
|
|
|
This notebook builds, trains, and runs a **small Transformer-based language model (mini GPT)** on a movie scripts dataset. |
|
|
Written for someone who knows **basic ML/DL** but is new to **LLMs**. |
|
|
|
|
|
--- |
|
|
|
|
|
## 1. Dataset & Preprocessing |
|
|
|
|
|
```python |
|
|
from datasets import load_dataset |
|
|
import tiktoken, numpy as np |
|
|
|
|
|
# Load dataset |
|
|
ds = load_dataset("IsmaelMousa/movies") |
|
|
|
|
|
# Split into train/val |
|
|
ds = ds['train'].train_test_split(test_size=0.1, seed=42) |
|
|
|
|
|
# Tokenizer (GPT-2) |
|
|
enc = tiktoken.get_encoding("gpt2") |
|
|
|
|
|
def process(example): |
|
|
ids = enc.encode_ordinary(example['Script']) |
|
|
return {'ids': ids, 'len': len(ids)} |
|
|
|
|
|
# Tokenize |
|
|
tokenized = ds.map(process, remove_columns=['Name','Script']) |
|
|
``` |
|
|
|
|
|
🔹 Dataset = movie scripts → tokenized into IDs → saved in `.bin` files for fast training. |
|
|
|
|
|
--- |
|
|
|
|
|
## 2. Create Input-Output Batches |
|
|
|
|
|
The model trains on fixed-length chunks (`block_size`) of tokens. |
|
|
Each batch contains input `X` and target `Y` sequences, where `Y` is shifted by 1 (next-token labels). |
|
|
|
|
|
```python |
|
|
def get_batch(split): |
|
|
data = train_data if split == 'train' else val_data |
|
|
ix = torch.randint(len(data) - block_size, (batch_size,)) |
|
|
x = torch.stack([torch.from_numpy(data[i:i+block_size].astype(np.int64)) for i in ix]) |
|
|
y = torch.stack([torch.from_numpy(data[i+1:i+block_size+1].astype(np.int64)) for i in ix]) |
|
|
return x.to(device), y.to(device) |
|
|
``` |
|
|
|
|
|
🔹 This is how we feed training data: **chunks of movie script → model learns to predict next token**. |
|
|
|
|
|
--- |
|
|
|
|
|
## 3. Model Architecture |
|
|
|
|
|
The model is a **stack of Transformer blocks**, similar to GPT-2. |
|
|
|
|
|
### (a) LayerNorm |
|
|
```python |
|
|
class LayerNorm(nn.Module): |
|
|
def __init__(self, ndim, bias): |
|
|
super().__init__() |
|
|
self.weight = nn.Parameter(torch.ones(ndim)) |
|
|
self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None |
|
|
def forward(self, x): |
|
|
return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5) |
|
|
``` |
|
|
- Normalizes features → stabilizes training. |
|
|
- Like BatchNorm, but per token, not per batch. |
|
|
|
|
|
--- |
|
|
|
|
|
### (b) Causal Self-Attention |
|
|
```python |
|
|
class CausalSelfAttention(nn.Module): |
|
|
def __init__(self, config): |
|
|
super().__init__() |
|
|
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) # QKV |
|
|
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) |
|
|
self.n_head = config.n_head |
|
|
self.n_embd = config.n_embd |
|
|
|
|
|
def forward(self, x): |
|
|
B, T, C = x.size() |
|
|
q, k, v = self.c_attn(x).split(self.n_embd, dim=2) |
|
|
# Reshape into multi-heads |
|
|
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
|
|
|
# Masked self-attention (causal: no peeking forward) |
|
|
att = (q @ k.transpose(-2, -1)) / (C // self.n_head)**0.5 |
|
|
mask = torch.tril(torch.ones(T, T, device=x.device)) |
|
|
att = att.masked_fill(mask == 0, float('-inf')) |
|
|
att = F.softmax(att, dim=-1) |
|
|
y = att @ v |
|
|
|
|
|
# Recombine heads |
|
|
y = y.transpose(1, 2).contiguous().view(B, T, C) |
|
|
return self.c_proj(y) |
|
|
``` |
|
|
- Lets each token "attend" to previous tokens. |
|
|
- Causal masking ensures left-to-right generation. |
|
|
|
|
|
--- |
|
|
|
|
|
### (c) MLP |
|
|
```python |
|
|
class MLP(nn.Module): |
|
|
def __init__(self, config): |
|
|
super().__init__() |
|
|
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd) |
|
|
self.gelu = nn.GELU() |
|
|
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd) |
|
|
|
|
|
def forward(self, x): |
|
|
return self.c_proj(self.gelu(self.c_fc(x))) |
|
|
``` |
|
|
- Expands hidden dim by 4x, then projects back. |
|
|
- Adds non-linear transformation. |
|
|
|
|
|
--- |
|
|
|
|
|
### (d) Transformer Block |
|
|
```python |
|
|
class Block(nn.Module): |
|
|
def __init__(self, config): |
|
|
super().__init__() |
|
|
self.ln1 = LayerNorm(config.n_embd, config.bias) |
|
|
self.attn = CausalSelfAttention(config) |
|
|
self.ln2 = LayerNorm(config.n_embd, config.bias) |
|
|
self.mlp = MLP(config) |
|
|
|
|
|
def forward(self, x): |
|
|
x = x + self.attn(self.ln1(x)) # Residual |
|
|
x = x + self.mlp(self.ln2(x)) # Residual |
|
|
return x |
|
|
``` |
|
|
- Core Transformer block = `[Norm → Attention → Residual → Norm → MLP → Residual]`. |
|
|
|
|
|
--- |
|
|
|
|
|
### (e) GPT Model |
|
|
```python |
|
|
class GPT(nn.Module): |
|
|
def __init__(self, config): |
|
|
super().__init__() |
|
|
self.transformer = nn.ModuleDict(dict( |
|
|
wte = nn.Embedding(config.vocab_size, config.n_embd), # token embedding |
|
|
wpe = nn.Embedding(config.block_size, config.n_embd), # position embedding |
|
|
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), |
|
|
ln_f = LayerNorm(config.n_embd, config.bias), |
|
|
)) |
|
|
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
|
|
self.transformer.wte.weight = self.lm_head.weight # weight tying |
|
|
|
|
|
def forward(self, idx, targets=None): |
|
|
b, t = idx.size() |
|
|
tok_emb = self.transformer.wte(idx) |
|
|
pos_emb = self.transformer.wpe(torch.arange(0, t, device=idx.device)) |
|
|
x = tok_emb + pos_emb |
|
|
for block in self.transformer.h: |
|
|
x = block(x) |
|
|
x = self.transformer.ln_f(x) |
|
|
logits = self.lm_head(x) |
|
|
|
|
|
if targets is None: |
|
|
return logits, None |
|
|
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) |
|
|
return logits, loss |
|
|
``` |
|
|
- Input tokens → embeddings + positional encoding → Transformer blocks → logits over vocab. |
|
|
- If `targets` provided → compute cross-entropy loss. |
|
|
- Otherwise → just output logits for generation. |
|
|
|
|
|
--- |
|
|
|
|
|
### (f) Generation |
|
|
```python |
|
|
@torch.no_grad() |
|
|
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): |
|
|
for _ in range(max_new_tokens): |
|
|
idx_cond = idx[:, -self.config.block_size:] |
|
|
logits, _ = self(idx_cond) |
|
|
logits = logits[:, -1, :] / temperature |
|
|
if top_k is not None: |
|
|
v, _ = torch.topk(logits, top_k) |
|
|
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 |
|
|
``` |
|
|
- Autoregressively generates tokens. |
|
|
- Uses `temperature` (randomness) and `top_k` (restricts to top-k likely tokens). |
|
|
|
|
|
--- |
|
|
|
|
|
## 4. Training |
|
|
|
|
|
- **Loss**: Cross-Entropy (predict next token). |
|
|
- **Optimizer**: AdamW (with tuned betas, weight decay). |
|
|
- **Scheduler**: Warmup + Cosine Decay. |
|
|
- **Mixed Precision + Gradient Accumulation** for efficiency. |
|
|
|
|
|
--- |
|
|
|
|
|
## 5. Monitoring |
|
|
|
|
|
```python |
|
|
plt.plot(train_loss_list, 'g', label='train_loss') |
|
|
plt.plot(validation_loss_list, 'r', label='validation_loss') |
|
|
plt.xlabel("Steps - Every 100 epochs") |
|
|
plt.ylabel("Loss") |
|
|
plt.legend() |
|
|
plt.show() |
|
|
``` |
|
|
- Green = training loss, Red = validation loss. |
|
|
- Watch for overfitting / underfitting. |
|
|
|
|
|
--- |
|
|
|
|
|
## 📊 Training Metrics |
|
|
|
|
|
| Epoch | Train Loss | Val Loss | Perplexity | |
|
|
|-------|------------|----------|------------| |
|
|
| 500 | 6.0358 | 6.0601 | 430.1 | |
|
|
| 1000 | 5.0690 | 5.1143 | 166.0 | |
|
|
| 1500 | 4.3162 | 4.3407 | 76.7 | |
|
|
| 2000 | 3.5948 | 3.6099 | 36.9 | |
|
|
| 2500 | 3.0460 | 3.0569 | 21.3 | |
|
|
| 3000 | 2.7518 | 2.7398 | 15.5 | |
|
|
| 3500 | 2.5606 | 2.5574 | 12.9 | |
|
|
| 4000 | 2.4583 | 2.4691 | 11.8 | |
|
|
| 4500 | 2.3943 | 2.3969 | 11.0 | |
|
|
| 5000 | 2.3428 | 2.3513 | 10.5 | |
|
|
| 6000 | 2.2141 | 2.2155 | 9.17 | |
|
|
| 7000 | 2.1389 | 2.1577 | 8.65 | |
|
|
| 8000 | 2.0570 | 2.0703 | 7.93 | |
|
|
| 9000 | 2.0062 | 2.0210 | 7.55 | |
|
|
| 10000 | 1.9604 | 1.9715 | 7.18 | |
|
|
| 12000 | 1.8580 | 1.8924 | 6.64 | |
|
|
| 14000 | 1.7954 | 1.8284 | 6.23 | |
|
|
| 16000 | 1.7369 | 1.7937 | 5.95 | |
|
|
| 18000 | 1.6901 | 1.7314 | 5.65 | |
|
|
| 19500 | 1.6594 | 1.7216 | 5.60 | |
|
|
|
|
|
📉 Validation loss steadily decreases, and **perplexity drops from ~430 → ~5.6** over training. |
|
|
|
|
|
## 6. Inference |
|
|
|
|
|
```python |
|
|
# Load best model |
|
|
model = GPT(config) |
|
|
model.load_state_dict(torch.load("best_model_params.pt", map_location=device)) |
|
|
model.eval() |
|
|
|
|
|
# Prompt |
|
|
sentence = "Write a Tarantino-style diner scene with two strangers..." |
|
|
context = torch.tensor(enc.encode_ordinary(sentence)).unsqueeze(0).to(device) |
|
|
|
|
|
# Generate (recommended shorter length) |
|
|
y = model.generate(context, max_new_tokens=300, temperature=0.8, top_k=50) |
|
|
print(enc.decode(y[0].tolist())) |
|
|
``` |
|
|
|
|
|
⚠️ Note: In the notebook, `max_new_tokens=5000` was used, which may be excessive. |
|
|
For practical testing, use **200–500 tokens**. |
|
|
|
|
|
--- |
|
|
|
|
|
## ✅ Summary |
|
|
|
|
|
- **Architecture**: GPT-like Transformer (attention + MLP blocks). |
|
|
- **Training**: Next-token prediction with AdamW + LR scheduling. |
|
|
- **Evaluation**: Loss curves (train vs val). |
|
|
- **Inference**: Autoregressive generation with temperature & top-k control. |
|
|
|
|
|
This is essentially a **mini GPT-2 clone**, scaled down for small datasets like movie scripts. |
