Metadata-Version: 2.4
Name: transformer-toolkit
Version: 0.0.16
Summary: Minimal, modular transformer library for training your own LLM
Author: Govind Barbade
Project-URL: Homepage, https://github.com/govindbarbade/transformer-toolkit
Project-URL: Repository, https://github.com/govindbarbade/transformer-toolkit
Keywords: transformer,llm,deep learning,nlp,pytorch
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.0.0
Requires-Dist: pydantic>=2.0.0
Provides-Extra: tokenizers
Requires-Dist: tokenizers>=0.15.0; extra == "tokenizers"
Provides-Extra: hf
Requires-Dist: transformers>=4.35.0; extra == "hf"
Requires-Dist: huggingface_hub>=0.20.0; extra == "hf"
Requires-Dist: datasets>=2.14.0; extra == "hf"
Provides-Extra: all
Requires-Dist: transformer-toolkit[hf,tokenizers]; extra == "all"

# transformer-toolkit

A modular, from-scratch transformer library for training and experimenting with modern LLM architectures. Swap attention types, positional encodings, FFN variants, and normalization — all from a single config object.

---

## Installation

```bash
pip install transformer-toolkit
```

---

## Quick Start

```python
import torch
from transformer_toolkit.model import Transformer, TransformerConfig
from transformer_toolkit.c_tokenizers import RustBPETokenizer
from transformer_toolkit.dataloader import DataConfig, from_binary, save_binary
from transformer_toolkit.trainer import Trainer, TrainConfig

# tokenizer
tok = RustBPETokenizer()
tok.train(open("data.txt", encoding="utf-8").readlines(), vocab_size=8000)
tok.save("tokenizer.json")

# data
save_binary(tok.encode(open("data.txt", encoding="utf-8").read()), "data.bin")
train_dl, val_dl = from_binary("data.bin", DataConfig(seq_len=128, batch_size=32))

# model
model = Transformer(TransformerConfig(
    vocab_size = tok.vocab_size,
    dim        = 512,
    n_layers   = 8,
    n_heads    = 8,
    pos_enc    = "rope",
)).to("cuda")

# train
trainer = Trainer(model, train_dl, val_dl, tok.vocab_size, TrainConfig(max_steps=3000))
trainer.train()
```

---

## Table of Contents

- [Model](#model)
- [Attention](#attention)
- [Feed-Forward Networks](#feed-forward-networks)
- [Positional Encodings](#positional-encodings)
- [Normalization](#normalization)
- [Dataloader](#dataloader)
- [Tokenizers](#tokenizers)
- [Trainer](#trainer)
- [HuggingFace Hub](#huggingface-hub)
- [Generation](#generation)
- [Full Examples](#full-examples)

---

## Model

### `TransformerConfig`

All architecture decisions live in one dataclass. Pass it to `Transformer()`.

```python
from transformer_toolkit.model import TransformerConfig

cfg = TransformerConfig(
    # ── core ──────────────────────────────────────
    vocab_size = 32000,      # tokenizer vocabulary size
    dim        = 512,        # model embedding dimension
    n_layers   = 8,          # number of transformer blocks
    n_heads    = 8,          # attention heads
    max_seq    = 2048,       # maximum sequence length

    # ── attention ─────────────────────────────────
    attn       = "gqa",      # "mha" | "gqa" | "mqa" | "flash" | "mla"
    n_kv_heads = 4,          # gqa only — number of key/value heads
    latent_dim = 64,         # mla only — latent compression dim

    # ── feed-forward ──────────────────────────────
    ffn        = "swiglu",   # "ffn" | "swiglu" | "moe"
    hidden_dim = 2048,       # FFN inner dimension (default: dim * 4)
    n_experts  = 8,          # moe only — number of experts
    top_k      = 2,          # moe only — experts activated per token
    moe_aux_weight = 0.01,   # moe load-balancing loss coefficient

    # ── normalization ─────────────────────────────
    norm       = "rmsnorm",  # "rmsnorm" | "layernorm"
    eps        = 1e-6,

    # ── positional encoding ───────────────────────
    pos_enc    = "rope",     # "rope" | "sinusoidal" | "learned" | "alibi" | "none"

    # ── regularisation ────────────────────────────
    dropout    = 0.1,

    # ── output ────────────────────────────────────
    tie_weights = True,      # share embedding and output projection weights
)
```

### `Transformer`

```python
from transformer_toolkit.model import Transformer

model = Transformer(cfg).to("cuda")

print(model.n_params())   # "30.21M"

# forward pass — returns (logits, aux_loss)
# aux_loss is non-zero only for MoE; always add it to your training loss
logits, aux_loss = model(tokens)       # tokens: [B, T]  →  logits: [B, T, vocab_size]

# generation
output = model.generate(
    tokens      = prompt_tokens,   # [B, T]
    max_new     = 200,
    temperature = 0.8,
    top_k       = 40,
)
```

> **Weight tying** — when `tie_weights=True` (default), the embedding and output projection share the same weight matrix. Use `model.state_dict_for_save()` instead of `model.state_dict()` when saving checkpoints, and `model.load_state_dict_with_tie()` when loading, to correctly preserve the tie across save/load cycles.

---

## Attention

Five attention variants, all swappable via `TransformerConfig.attn`.

| Value | Class | Used in |
|-------|-------|---------|
| `"mha"` | `MultiHeadAttention` | Original Transformer, BERT, GPT-2 |
| `"gqa"` | `GroupedQueryAttention` | LLaMA 3, Mistral |
| `"mqa"` | `MultiQueryAttention` | Falcon, early Gemini |
| `"flash"` | `FlashAttention` | Any model on PyTorch ≥ 2.0 |
| `"mla"` | `MLAttention` | DeepSeek-V2/V3 |

**RoPE** is applied inside attention to `q` and `k` after head-splitting — not to the residual stream. It is instantiated once and shared across all layers.

**ALiBi** bias is computed once per forward pass and passed as an additive mask to every block.

### Example — switch to Flash Attention

```python
cfg = TransformerConfig(
    dim     = 512,
    n_heads = 8,
    attn    = "flash",   # uses torch.nn.functional.scaled_dot_product_attention
)
```

### Example — Grouped Query Attention (LLaMA-style)

```python
cfg = TransformerConfig(
    dim        = 512,
    n_heads    = 8,
    attn       = "gqa",
    n_kv_heads = 2,   # 4 query heads share each kv head → 4x kv cache reduction
)
```

---

## Feed-Forward Networks

| Value | Class | Used in |
|-------|-------|---------|
| `"ffn"` | `FFN` | Original Transformer, BERT |
| `"swiglu"` | `SwiGLU` | LLaMA, Mistral, PaLM |
| `"moe"` | `MoE` | Mixtral, GPT-4 (rumoured) |

### MoE — Mixture of Experts

When using `ffn="moe"`, the model forward pass returns an auxiliary load-balancing loss. You **must** add it to your training loss or all tokens will collapse onto 1–2 experts within a few hundred steps.

```python
cfg = TransformerConfig(
    ffn            = "moe",
    n_experts      = 8,
    top_k          = 2,
    moe_aux_weight = 0.01,   # weight of the load-balancing term
)

logits, aux_loss = model(tokens)
ce_loss  = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
loss     = ce_loss + aux_loss   # aux_loss is 0.0 for non-MoE models
```

---

## Positional Encodings

| Value | Applied | Notes |
|-------|---------|-------|
| `"rope"` | Inside attention, on q and k | LLaMA, Mistral, Qwen |
| `"sinusoidal"` | Residual stream before blocks | Original Transformer |
| `"learned"` | Residual stream before blocks | BERT, GPT-2 |
| `"alibi"` | Additive bias on attention scores | ALiBi paper |
| `"none"` | Not applied | Bare model, no position info |

Each encoding applies exactly once in exactly one place — there is no double-application.

---

## Normalization

| Value | Class | Notes |
|-------|-------|-------|
| `"rmsnorm"` | `RMSNorm` | LLaMA, Mistral, Qwen — no mean subtraction, no bias |
| `"layernorm"` | `LayerNorm` | BERT, GPT-2 — classic with bias |

---

## Dataloader

### `DataConfig`

```python
from transformer_toolkit.dataloader import DataConfig

cfg = DataConfig(
    seq_len     = 128,    # sequence length fed to the model
    batch_size  = 32,     # samples per batch
    split       = 0.9,    # fraction of data used for training
    stride      = None,   # None = non-overlapping windows (recommended)
                          # set stride < seq_len for overlapping windows
    shuffle     = True,
    num_workers = 4,
    pin_memory  = True,
    debug       = False,  # print sample decoding before training
    debug_n     = 3,      # number of samples to show when debug=True
)
```

> **stride** — the default `stride=None` (equivalent to `stride=seq_len`) produces non-overlapping windows. For a 1.86M token dataset with `seq_len=128` this gives ~14,600 clean distinct samples. Setting `stride=1` gives 1.86M heavily-overlapping samples and causes rapid overfitting on small datasets.

### Loading from a binary file

```python
from transformer_toolkit.dataloader import save_binary, from_binary

# tokenize once and save
save_binary(tok.encode(text), "data.bin")

# load — supports both raw uint16 binary and .npy
train_dl, val_dl = from_binary("data.bin", cfg, tokenizer=tok)

# optionally save the splits as .npy for memmap reuse on future runs
train_dl, val_dl = from_binary(
    "data.bin", cfg,
    train_path = "train.npy",
    val_path   = "val.npy",
    tokenizer  = tok,
)
```

### Memmap — loading pre-split .npy files (zero RAM)

On the second run, load the pre-split files directly. The token file stays on disk — only the pages you access load into RAM. Scales to 100GB+ datasets.

```python
from transformer_toolkit.dataloader import from_npy_split

train_dl, val_dl = from_npy_split("train.npy", "val.npy", cfg, tokenizer=tok)
```

### Loading from text files

```python
from transformer_toolkit.dataloader import from_files

train_dl, val_dl = from_files(
    paths      = ["data1.txt", "data2.txt"],
    tokenizer  = tok,
    cfg        = cfg,
    train_path = "train.npy",   # optional — save for future memmap reuse
    val_path   = "val.npy",
    bos_id     = tok.bos_id,    # optional — wrap documents with BOS/EOS
    eos_id     = tok.eos_id,
)
```

### Loading from HuggingFace

```python
from transformer_toolkit.dataloader import from_hf

# streaming — no full download, infinite dataset
cfg_stream = DataConfig(seq_len=512, batch_size=16, streaming=True)
train_dl, val_dl = from_hf("roneneldan/TinyStories", tok, cfg_stream)

# in-memory — downloads fully, then splits and optionally saves as .npy
train_dl, val_dl = from_hf(
    dataset_name = "roneneldan/TinyStories",
    tokenizer    = tok,
    cfg          = cfg,
    text_col     = "text",
    bos_id       = 1,
    eos_id        = 2,
    train_path   = "train.npy",
    val_path     = "val.npy",
)
```

### Debug mode

```python
cfg = DataConfig(seq_len=128, batch_size=32, debug=True, debug_n=3)
train_dl, val_dl = from_binary("data.bin", cfg, tokenizer=tok)
```

Prints before training starts:

```
──────────────────────────────────────────────────────────────
  🔍 Debug samples (train)
  seq_len=128  stride=128  batch_size=32
──────────────────────────────────────────────────────────────

  sample 1
  x ids : [23, 451, 12, 8, 1203, 44, 91 ...] ... +121
  y ids : [451, 12, 8, 1203, 44, 91, 7  ...] ... +121
  x text: 'ROMEO:\nBut soft, what light through yonder window...'
  y text: '\nBut soft, what light through yonder window breaks'
  ✓  x/y alignment correct (y = x shifted by 1)
```

---

## Tokenizers

Three tokenizers available, all sharing the same interface.

```python
from transformer_toolkit.c_tokenizers import (
    ByteLevelTokenizer,
    RustBPETokenizer,
    HFTokenizer,
)
```

### `ByteLevelTokenizer`

Zero dependencies. Every byte is a token (vocab size = 256). Works on any language out of the box.

```python
tok = ByteLevelTokenizer()
ids = tok.encode("Hello world")   # [72, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100]
txt = tok.decode(ids)             # "Hello world"
```

### `RustBPETokenizer`

BPE tokenizer backed by HuggingFace's Rust `tokenizers` library. Trains ~100x faster than pure Python BPE.

```python
pip install tokenizers
```

```python
tok = RustBPETokenizer()
tok.train(open("data.txt").readlines(), vocab_size=8000)
tok.save("tokenizer.json")

# later
tok.load("tokenizer.json")
ids = tok.encode("Hello world")
txt = tok.decode(ids)
print(tok.vocab_size)   # 8000
```

### `HFTokenizer`

Thin wrapper around any HuggingFace tokenizer.

```python
pip install transformers
```

```python
tok = HFTokenizer("gpt2")
ids = tok.encode("Hello world")
txt = tok.decode(ids)
print(tok.vocab_size)   # 50257
```

---

## Trainer

### `TrainConfig`

```python
from transformer_toolkit.trainer import TrainConfig

cfg = TrainConfig(
    # ── steps ─────────────────────────────────────
    max_steps        = 10000,   # total optimizer steps
    eval_every       = 500,     # run validation every N steps
    save_every       = 1000,    # save step_N.pt every N steps
    log_every        = 50,      # print loss to console every N steps
    interruptible    = True,    # Ctrl+C saves a clean checkpoint

    # ── optimiser ─────────────────────────────────
    lr               = 3e-4,   # peak learning rate
    min_lr           = 3e-5,   # floor lr after cosine decay
    weight_decay     = 0.1,
    beta1            = 0.9,
    beta2            = 0.95,
    grad_clip        = 1.0,    # max gradient norm

    # ── lr schedule ───────────────────────────────
    warmup_steps     = 200,    # linear warmup from 0 to lr

    # ── efficiency ────────────────────────────────
    grad_accum_steps = 4,      # effective batch = batch_size × grad_accum_steps
    mixed_precision  = True,   # bf16/fp16 on CUDA, fp32 on CPU
    grad_checkpoint  = False,  # recompute activations to save VRAM (~20% slower)

    # ── checkpoints ───────────────────────────────
    ckpt_dir         = "checkpoints",
    save_best        = True,   # save best.pt when val loss improves
    save_step_ckpts  = True,   # save step_N.pt every save_every steps

    # ── huggingface hub ───────────────────────────
    hf_repo          = "username/my-model",
    hf_private       = True,
    hf_push_best     = True,
    hf_push_every_n  = False,
    hf_push_end      = True,
    hf_push_on_pause = True,
)
```

### `Trainer`

```python
from transformer_toolkit.trainer import Trainer

trainer = Trainer(
    model      = model,
    train_dl   = train_dl,
    val_dl     = val_dl,
    vocab_size = tok.vocab_size,
    cfg        = cfg_train,
    tokenizer  = tok,          # optional — used for HF hub uploads
)

trainer.train()

# resume from a checkpoint
trainer.train(resume_from="checkpoints/step_2000.pt")
```

Training output:

```
──────────────────────────────────────────────────────────────
  ⚡ Transformer Toolkit Trainer
  steps=3000  lr=0.0003  warmup=200  accum=4
  mixed_precision=True  grad_clip=1.0
──────────────────────────────────────────────────────────────

  step    100/3000  ████████░░░░░░░░░░░░░░░░░░░░  loss 3.1423  lr 1.5e-04  eta 4m
  step    200/3000  ████████████░░░░░░░░░░░░░░░░  loss 2.8901  lr 3.0e-04  eta 3m

  ● eval  step 300  val_loss 2.7130  ppl 15.07  ▼0.1823  ★ best
```

---

## HuggingFace Hub

### Login

```python
from transformer_toolkit.hf_hub import login

login(token="hf_your_token_here")
# or
login(username="you", password="your_password")
```

### Push to Hub

```python
from transformer_toolkit.hf_hub import push_to_hub

push_to_hub(
    repo_id  = "username/my-model",
    model    = model,
    cfg      = cfg_model,
    tokenizer = tok,
    metrics  = {"val_loss": 1.83, "perplexity": 6.23},
    step     = 3000,
    private  = True,
)
```

### Pull from Hub

```python
from transformer_toolkit.hf_hub import pull_from_hub

pull_from_hub("username/my-model", save_dir="checkpoints")
# downloads model.pt, tokenizer.json, config.json, metrics.json
```

---

## Generation

```python
from transformer_toolkit.model import Transformer, TransformerConfig
from transformer_toolkit.c_tokenizers import RustBPETokenizer
from transformer_toolkit.trainer import load_ckpt
import torch

DEVICE = torch.device("cuda")

tok = RustBPETokenizer()
tok.load("tokenizer.json")

cfg = TransformerConfig(
    vocab_size = tok.vocab_size,
    dim        = 512,
    n_layers   = 8,
    n_heads    = 8,
    attn       = "gqa",
    n_kv_heads = 4,
    ffn        = "swiglu",
    hidden_dim = 2048,
    norm       = "rmsnorm",
    pos_enc    = "rope",
    dropout    = 0.0,   # always 0 at inference
)
model = Transformer(cfg).to(DEVICE)
load_ckpt("checkpoints/best.pt", model)
model.eval()

def generate(prompt, max_new=200, temperature=0.8, top_k=40):
    ids    = tok.encode(prompt)
    tokens = torch.tensor([ids], dtype=torch.long, device=DEVICE)
    out    = model.generate(tokens, max_new=max_new,
                             temperature=temperature, top_k=top_k)
    return tok.decode(out[0].tolist())

print(generate("ROMEO:"))
```

**Generation parameters:**

| Parameter | Effect | Recommended |
|-----------|--------|-------------|
| `temperature` | Higher = more random, lower = more repetitive | 0.7 – 1.0 |
| `top_k` | Only sample from top-k tokens | 20 – 50 |
| `max_new` | Number of new tokens to generate | 100 – 500 |

---

## Full Examples

### Small model — Shakespeare (< 2M tokens, any GPU)

```python
import torch, os
from transformer_toolkit.model import Transformer, TransformerConfig
from transformer_toolkit.c_tokenizers import RustBPETokenizer
from transformer_toolkit.dataloader import DataConfig, from_binary, save_binary
from transformer_toolkit.trainer import Trainer, TrainConfig

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tok = RustBPETokenizer()
tok.train(open("shakespeare.txt", encoding="utf-8").readlines(), vocab_size=4000)
tok.save("tokenizer.json")

if not os.path.exists("data.bin"):
    save_binary(tok.encode(open("shakespeare.txt", encoding="utf-8").read()), "data.bin")

train_dl, val_dl = from_binary("data.bin", DataConfig(
    seq_len=128, batch_size=32, split=0.9, stride=None
), tokenizer=tok)

model = Transformer(TransformerConfig(
    vocab_size = tok.vocab_size,
    dim        = 384,
    n_layers   = 6,
    n_heads    = 6,
    n_kv_heads = 3,
    attn       = "gqa",
    ffn        = "swiglu",
    hidden_dim = 1536,
    norm       = "rmsnorm",
    pos_enc    = "rope",
    dropout    = 0.1,
)).to(DEVICE)
print(model.n_params())   # ~15M

trainer = Trainer(model, train_dl, val_dl, tok.vocab_size, TrainConfig(
    max_steps        = 3000,
    warmup_steps     = 200,
    eval_every       = 300,
    lr               = 3e-4,
    grad_accum_steps = 4,
    mixed_precision  = True,
    save_best        = True,
    save_step_ckpts  = True,
))
trainer.train()
```

### Large dataset — HuggingFace streaming

```python
from transformer_toolkit.dataloader import DataConfig, from_hf, from_npy_split
from transformer_toolkit.c_tokenizers import HFTokenizer

tok = HFTokenizer("HuggingFaceTB/SmolLM-135M")

cfg = DataConfig(seq_len=512, batch_size=16, stride=None, num_workers=4)

# first run — tokenizes and saves splits as memmap .npy
train_dl, val_dl = from_hf(
    dataset_name = "roneneldan/TinyStories",
    tokenizer    = tok,
    cfg          = cfg,
    bos_id       = tok._tok.bos_token_id,
    eos_id       = tok._tok.eos_token_id,
    train_path   = "train.npy",
    val_path     = "val.npy",
)

# second+ runs — zero RAM, instant load
train_dl, val_dl = from_npy_split("train.npy", "val.npy", cfg, tokenizer=tok)
```

### MoE model

```python
model = Transformer(TransformerConfig(
    vocab_size     = tok.vocab_size,
    dim            = 512,
    n_layers       = 8,
    n_heads        = 8,
    attn           = "flash",
    ffn            = "moe",
    n_experts      = 8,
    top_k          = 2,
    moe_aux_weight = 0.01,
    pos_enc        = "rope",
    dropout        = 0.1,
)).to("cuda")

# model.forward() returns (logits, aux_loss)
# TrainConfig handles this automatically — no changes needed in training code
trainer = Trainer(model, train_dl, val_dl, tok.vocab_size, TrainConfig(
    max_steps = 5000,
    lr        = 3e-4,
))
trainer.train()
```

---

## Architecture Reference

```
Input tokens [B, T]
      │
      ▼
Embedding [B, T, dim]  +  positional encoding (sinusoidal / learned)
      │
      ▼  × n_layers
┌─────────────────────────────────┐
│  Norm → Attention (+ RoPE/ALiBi)│
│  Residual                       │
│  Norm → FFN / SwiGLU / MoE      │
│  Residual                       │
└─────────────────────────────────┘
      │
      ▼
Final Norm → Linear head [B, T, vocab_size]
```

---

## Requirements

- Python ≥ 3.10
- PyTorch ≥ 2.0
- `tokenizers` — for `RustBPETokenizer` (`pip install tokenizers`)
- `transformers` — for `HFTokenizer` (`pip install transformers`)
- `datasets` — for `from_hf()` (`pip install datasets`)
- `huggingface_hub` — for hub push/pull (`pip install huggingface_hub`)
- `pydantic` — for `TrainConfig` validation (`pip install pydantic`)

---

## License

MIT
