Add CPU-trained tiny character LLM

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+.DS_Store

README.md

@@ -0,0 +1,80 @@
+---
+license: mit
+library_name: pytorch
+tags:
+- tinyllm
+- character-language-model
+- gpt
+- cpu
+- educational
+- pytorch
+pipeline_tag: text-generation
+---
+
+# tinyllm-cpu-char
+
+A tiny CPU-trained character-level GPT-style language model created as an educational end-to-end LLM experiment.
+
+This is **not** a general chatbot. It is a deliberately small model trained on a tiny toy corpus to demonstrate the full loop:
+
+1. build a vocabulary
+2. train a causal transformer
+3. save a checkpoint
+4. generate text from learned character patterns
+
+## Model details
+
+- Architecture: tiny GPT-style causal transformer
+- Tokenization: character-level
+- Parameters: 106,688
+- Training device: CPU
+- Training steps: 3,000
+- Dataset: tiny included toy corpus, `data/tiny_corpus.txt`
+- Checkpoint: `checkpoints/tinyllm_overfit_3k.pt`
+
+Final training run:
+
+```text
+step 0:    train 3.6967, val 3.6921
+step 3000: train 0.1903, val 3.7094
+```
+
+The rising validation loss is expected here: this checkpoint intentionally overfits the tiny corpus to prove the training loop works.
+
+## Quick inference
+
+```bash
+pip install torch huggingface_hub
+python hf_infer.py --repo-id YOUR_USERNAME/tinyllm-cpu-char --prompt "A seed" --tokens 300
+```
+
+For local inference after cloning:
+
+```bash
+python sample.py --ckpt checkpoints/tinyllm_overfit_3k.pt --prompt "The little machine" --tokens 300
+```
+
+Example output after overfitting:
+
+```text
+A seed does not pretend to be a forest.
+A sed shows that a forest is posssible.
+
+The student asked: can a tiny model think?
+The teacher answered: first let it predict...
+```
+
+## Training
+
+```bash
+python train.py --steps 3000 --eval-interval 500 --eval-iters 10 \
+  --batch-size 16 --n-embd 64 --block-size 64 \
+  --out checkpoints/tinyllm_overfit_3k.pt
+```
+
+## Limitations
+
+- Character-level only.
+- Trained on a tiny toy corpus.
+- Generates memorized/stylized snippets, not reliable knowledge.
+- Intended for learning and experimentation.

checkpoints/tinyllm_overfit_3k.json

@@ -0,0 +1,48 @@
+{
+  "params": 106688,
+  "chars": [
+    "\n",
+    " ",
+    ",",
+    ".",
+    ":",
+    "?",
+    "A",
+    "C",
+    "E",
+    "I",
+    "P",
+    "T",
+    "U",
+    "W",
+    "a",
+    "b",
+    "c",
+    "d",
+    "e",
+    "f",
+    "g",
+    "h",
+    "i",
+    "k",
+    "l",
+    "m",
+    "n",
+    "o",
+    "p",
+    "q",
+    "r",
+    "s",
+    "t",
+    "u",
+    "v",
+    "w",
+    "x",
+    "y",
+    "z"
+  ],
+  "last_loss": {
+    "train": 0.19029191136360168,
+    "val": 3.7094032764434814
+  }
+}

checkpoints/tinyllm_overfit_3k.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:068f3a2580674876837e22022acf895f2bc67cd0f3d30a3782802373a0e01b0d
+size 438789

data/tiny_corpus.txt

@@ -0,0 +1,34 @@
+The little machine listened to the rain and learned one letter at a time.
+A thought began as a spark, then became a pattern, then became a voice.
+The teacher said: do not fear small beginnings. A tiny model can still teach us the shape of learning.
+
+In the quiet workshop, tokens walked in a line.
+Each token looked back at the tokens before it and asked, what should come next?
+The answer was never magic. It was counting, guessing, correcting, and trying again.
+
+A fox told a crow: wisdom is not the size of the library, but the care of the attention.
+A crow replied: even a small mind can remember a melody if it hears the song often enough.
+The river laughed, because the river had trained every stone by repeating its lesson.
+
+We build tinyllm to understand language models from the inside.
+We train on CPU because patience is part of the experiment.
+We start with characters because characters are honest: small marks, simple rules, many possibilities.
+
+The model begins with random noise.
+Then loss falls a little.
+Then letters become syllables.
+Then syllables become words.
+Then words begin to imitate the books that fed them.
+
+This is not a giant assistant.
+This is a seed.
+A seed does not pretend to be a forest.
+A seed shows that a forest is possible.
+
+The student asked: can a tiny model think?
+The teacher answered: first let it predict. Then let us study what prediction teaches.
+
+Again the little machine listened.
+Again the optimizer stepped.
+Again the text became less strange.
+And in the warm hum of the CPU, the tiny language model began.

hf_infer.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+"""Run inference from a local clone or directly from a Hugging Face repo."""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+from huggingface_hub import hf_hub_download
+
+from model import TinyGPT, TinyGPTConfig
+
+
+def resolve_ckpt(args) -> Path:
+    if args.ckpt:
+        return Path(args.ckpt)
+    if args.repo_id:
+        return Path(
+            hf_hub_download(
+                repo_id=args.repo_id,
+                filename=args.filename,
+                revision=args.revision,
+            )
+        )
+    local = Path(args.filename)
+    if local.exists():
+        return local
+    raise SystemExit("Provide --ckpt for local checkpoint or --repo-id for Hugging Face download.")
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--repo-id", help="Hugging Face repo id, e.g. username/tinyllm-cpu-char")
+    p.add_argument("--revision", default="main")
+    p.add_argument("--filename", default="checkpoints/tinyllm_overfit_3k.pt")
+    p.add_argument("--ckpt", help="Local checkpoint path")
+    p.add_argument("--prompt", default="The little machine")
+    p.add_argument("--tokens", type=int, default=300)
+    p.add_argument("--temperature", type=float, default=0.7)
+    p.add_argument("--top-k", type=int, default=10)
+    args = p.parse_args()
+
+    ckpt = torch.load(resolve_ckpt(args), map_location="cpu")
+    cfg = TinyGPTConfig(**ckpt["config"])
+    model = TinyGPT(cfg)
+    model.load_state_dict(ckpt["model_state"])
+    model.eval()
+
+    stoi = ckpt["stoi"]
+    itos = {int(k): v for k, v in ckpt["itos"].items()}
+    prompt = "".join(ch for ch in args.prompt if ch in stoi) or "\n"
+    idx = torch.tensor([[stoi[ch] for ch in prompt]], dtype=torch.long)
+    out = model.generate(idx, max_new_tokens=args.tokens, temperature=args.temperature, top_k=args.top_k)
+    print("".join(itos[int(i)] for i in out[0]))
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,127 @@
+"""A deliberately tiny GPT-style language model for CPU experiments."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class TinyGPTConfig:
+    def __init__(
+        self,
+        vocab_size: int,
+        block_size: int = 64,
+        n_layer: int = 2,
+        n_head: int = 2,
+        n_embd: int = 64,
+        dropout: float = 0.1,
+    ):
+        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
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, cfg: TinyGPTConfig):
+        super().__init__()
+        assert cfg.n_embd % cfg.n_head == 0
+        self.n_head = cfg.n_head
+        self.head_dim = cfg.n_embd // cfg.n_head
+        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd)
+        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd)
+        self.dropout = nn.Dropout(cfg.dropout)
+        self.register_buffer(
+            "mask",
+            torch.tril(torch.ones(cfg.block_size, cfg.block_size)).view(1, 1, cfg.block_size, cfg.block_size),
+            persistent=False,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, t, c = x.shape
+        q, k, v = self.qkv(x).split(c, dim=2)
+        q = q.view(b, t, self.n_head, self.head_dim).transpose(1, 2)
+        k = k.view(b, t, self.n_head, self.head_dim).transpose(1, 2)
+        v = v.view(b, t, self.n_head, self.head_dim).transpose(1, 2)
+
+        att = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
+        att = att.masked_fill(self.mask[:, :, :t, :t] == 0, float("-inf"))
+        att = F.softmax(att, dim=-1)
+        att = self.dropout(att)
+        y = att @ v
+        y = y.transpose(1, 2).contiguous().view(b, t, c)
+        return self.dropout(self.proj(y))
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: TinyGPTConfig):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(cfg.n_embd)
+        self.attn = CausalSelfAttention(cfg)
+        self.ln2 = nn.LayerNorm(cfg.n_embd)
+        self.mlp = nn.Sequential(
+            nn.Linear(cfg.n_embd, 4 * cfg.n_embd),
+            nn.GELU(),
+            nn.Linear(4 * cfg.n_embd, cfg.n_embd),
+            nn.Dropout(cfg.dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class TinyGPT(nn.Module):
+    def __init__(self, cfg: TinyGPTConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.token_embedding = nn.Embedding(cfg.vocab_size, cfg.n_embd)
+        self.position_embedding = nn.Embedding(cfg.block_size, cfg.n_embd)
+        self.drop = nn.Dropout(cfg.dropout)
+        self.blocks = nn.Sequential(*[Block(cfg) for _ in range(cfg.n_layer)])
+        self.ln_f = nn.LayerNorm(cfg.n_embd)
+        self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
+
+        # Weight tying: common in GPT-style LMs.
+        self.head.weight = self.token_embedding.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
+        b, t = idx.shape
+        if t > self.cfg.block_size:
+            raise ValueError(f"sequence length {t} > block_size {self.cfg.block_size}")
+        pos = torch.arange(0, t, device=idx.device)
+        x = self.token_embedding(idx) + self.position_embedding(pos)[None, :, :]
+        x = self.drop(x)
+        x = self.blocks(x)
+        x = self.ln_f(x)
+        logits = self.head(x)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx: torch.Tensor, max_new_tokens: int, temperature: float = 0.8, top_k: int | None = None):
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.cfg.block_size :]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / max(temperature, 1e-6)
+            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)
+            next_idx = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, next_idx), dim=1)
+        return idx

requirements.txt

@@ -0,0 +1,2 @@
+torch>=2.0
+huggingface_hub>=0.23

sample.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+"""Generate text from a tinyllm checkpoint."""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+
+from model import TinyGPT, TinyGPTConfig
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", default="checkpoints/tinyllm.pt")
+    p.add_argument("--prompt", default="The")
+    p.add_argument("--tokens", type=int, default=300)
+    p.add_argument("--temperature", type=float, default=0.8)
+    p.add_argument("--top-k", type=int, default=20)
+    args = p.parse_args()
+
+    ckpt = torch.load(Path(args.ckpt), map_location="cpu")
+    cfg = TinyGPTConfig(**ckpt["config"])
+    model = TinyGPT(cfg)
+    model.load_state_dict(ckpt["model_state"])
+    model.eval()
+
+    stoi = ckpt["stoi"]
+    itos = {int(k): v for k, v in ckpt["itos"].items()}
+    safe_prompt = "".join(ch for ch in args.prompt if ch in stoi) or "\n"
+    idx = torch.tensor([[stoi[ch] for ch in safe_prompt]], dtype=torch.long)
+    out = model.generate(idx, max_new_tokens=args.tokens, temperature=args.temperature, top_k=args.top_k)
+    text = "".join(itos[int(i)] for i in out[0])
+    print(text)
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+"""Train a tiny character-level GPT on CPU.
+
+This is intentionally small and educational, not production-grade.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import torch
+
+from model import TinyGPT, TinyGPTConfig
+
+
+def build_vocab(text: str):
+    chars = sorted(set(text))
+    stoi = {ch: i for i, ch in enumerate(chars)}
+    itos = {i: ch for ch, i in stoi.items()}
+    return chars, stoi, itos
+
+
+def encode(text: str, stoi: dict[str, int]):
+    return [stoi[ch] for ch in text]
+
+
+def get_batch(data: torch.Tensor, block_size: int, batch_size: int, device: str):
+    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
+    x = torch.stack([data[i : i + block_size] for i in ix]).to(device)
+    y = torch.stack([data[i + 1 : i + block_size + 1] for i in ix]).to(device)
+    return x, y
+
+
+@torch.no_grad()
+def estimate_loss(model, train_data, val_data, block_size, batch_size, eval_iters, device):
+    out = {}
+    model.eval()
+    for split, data in [("train", train_data), ("val", val_data)]:
+        losses = torch.zeros(eval_iters)
+        for k in range(eval_iters):
+            xb, yb = get_batch(data, block_size, batch_size, device)
+            _, loss = model(xb, yb)
+            losses[k] = loss.item()
+        out[split] = losses.mean().item()
+    model.train()
+    return out
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data", default="data/tiny_corpus.txt")
+    p.add_argument("--out", default="checkpoints/tinyllm.pt")
+    p.add_argument("--steps", type=int, default=500)
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--block-size", type=int, default=64)
+    p.add_argument("--n-layer", type=int, default=2)
+    p.add_argument("--n-head", type=int, default=2)
+    p.add_argument("--n-embd", type=int, default=64)
+    p.add_argument("--lr", type=float, default=3e-4)
+    p.add_argument("--eval-interval", type=int, default=100)
+    p.add_argument("--eval-iters", type=int, default=10)
+    p.add_argument("--seed", type=int, default=1337)
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    device = "cpu"
+
+    data_path = Path(args.data)
+    text = data_path.read_text(encoding="utf-8")
+    if len(text) < args.block_size + 2:
+        raise SystemExit("Dataset is too small for the chosen block size.")
+
+    chars, stoi, itos = build_vocab(text)
+    encoded = torch.tensor(encode(text, stoi), dtype=torch.long)
+    n = int(0.9 * len(encoded))
+    train_data = encoded[:n]
+    val_data = encoded[n:] if len(encoded[n:]) > args.block_size + 1 else encoded[:n]
+
+    cfg = TinyGPTConfig(
+        vocab_size=len(chars),
+        block_size=args.block_size,
+        n_layer=args.n_layer,
+        n_head=args.n_head,
+        n_embd=args.n_embd,
+        dropout=0.1,
+    )
+    model = TinyGPT(cfg).to(device)
+    params = sum(p.numel() for p in model.parameters())
+    print(f"chars={len(chars)} tokens={len(encoded)} params={params:,} device={device}")
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
+    start = time.time()
+    last_loss = None
+    for step in range(args.steps + 1):
+        if step % args.eval_interval == 0 or step == args.steps:
+            losses = estimate_loss(model, train_data, val_data, args.block_size, args.batch_size, args.eval_iters, device)
+            print(f"step {step:5d}: train {losses['train']:.4f}, val {losses['val']:.4f}")
+            last_loss = losses
+
+        xb, yb = get_batch(train_data, args.block_size, args.batch_size, device)
+        _, loss = model(xb, yb)
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+
+    out_path = Path(args.out)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    ckpt = {
+        "model_state": model.state_dict(),
+        "config": cfg.__dict__,
+        "stoi": stoi,
+        "itos": {str(k): v for k, v in itos.items()},
+        "train_args": vars(args),
+        "last_loss": last_loss,
+    }
+    torch.save(ckpt, out_path)
+    meta_path = out_path.with_suffix(".json")
+    meta_path.write_text(json.dumps({"params": params, "chars": chars, "last_loss": last_loss}, indent=2), encoding="utf-8")
+    print(f"saved {out_path} and {meta_path} in {time.time() - start:.1f}s")
+
+
+if __name__ == "__main__":
+    main()