| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """Pre-train and sample Eça, a small GPT trained from scratch on the works of Eça de Queirós. |
| |
| Byte-level BPE tokenizer trained on the corpus, pre-norm decoder-only transformer with tied |
| embeddings. Model quality is reported as bits per character on the last 2% of the text, so it |
| stays comparable across tokenizers. |
| """ |
|
|
| |
| |
|
|
| import argparse |
| import copy |
| import json |
| import math |
| import shutil |
| import time |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, snapshot_download |
| from huggingface_hub.errors import RepositoryNotFoundError |
| from safetensors.torch import load_file, save_file |
| from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers |
| from torch import nn |
|
|
| REPO = "jgalego/eca" |
| DATA = "jgalego/eca-queiros" |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
| TOP_K = 40 |
| HERE = Path(__file__).parent |
|
|
|
|
| @dataclass |
| class Config: |
| """Architecture.""" |
|
|
| vocab: int = 2048 |
| block: int = 256 |
| width: int = 256 |
| layers: int = 6 |
| heads: int = 4 |
| dropout: float = 0.1 |
|
|
|
|
| |
|
|
|
|
| def read_text(source): |
| """Text of a local file or of eça.txt in the dataset repo.""" |
| if Path(source).is_file(): |
| return Path(source).read_text(encoding="utf-8") |
| return Path(hf_hub_download(source, "eça.txt", repo_type="dataset")).read_text(encoding="utf-8") |
|
|
|
|
| def new_tokenizer(text, vocab): |
| """Train a byte-level BPE tokenizer on the corpus.""" |
| tokenizer = Tokenizer(models.BPE()) |
| tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| tokenizer.decoder = decoders.ByteLevel() |
| alphabet = pre_tokenizers.ByteLevel.alphabet() |
| trainer = trainers.BpeTrainer(vocab_size=vocab, initial_alphabet=alphabet) |
| tokenizer.train_from_iterator(text.splitlines(keepends=True), trainer) |
| return tokenizer |
|
|
|
|
| def prepare(args): |
| """Tokenizer, training tokens, and the validation tokens with their length in characters. |
| |
| The first 98% of the text is for training, the rest for validation. |
| """ |
| text = read_text(args.data) |
| tokenizer = new_tokenizer(text, args.vocab) |
| cut = text.index("\n", int(len(text) * 0.98)) + 1 |
| head, tail = (torch.tensor(tokenizer.encode(part).ids) for part in (text[:cut], text[cut:])) |
| return tokenizer, head, (tail, len(text) - cut) |
|
|
|
|
| |
|
|
|
|
| class Block(nn.Module): |
| """Pre-norm transformer block with causal self-attention.""" |
|
|
| def __init__(self, config): |
| super().__init__() |
| self.heads, self.dropout = config.heads, config.dropout |
| self.norm1 = nn.LayerNorm(config.width) |
| self.qkv = nn.Linear(config.width, 3 * config.width) |
| self.proj = nn.Linear(config.width, config.width) |
| self.norm2 = nn.LayerNorm(config.width) |
| self.mlp = nn.Sequential( |
| nn.Linear(config.width, 4 * config.width), |
| nn.GELU(), |
| nn.Linear(4 * config.width, config.width), |
| ) |
| self.drop = nn.Dropout(config.dropout) |
|
|
| def forward(self, x): |
| """Attention then MLP, each with a residual connection.""" |
| batch, length, width = x.shape |
| qkv = self.qkv(self.norm1(x)).view(batch, length, 3, self.heads, -1) |
| q, k, v = qkv.permute(2, 0, 3, 1, 4) |
| dropout = self.dropout if self.training else 0.0 |
| att = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=dropout) |
| x = x + self.drop(self.proj(att.transpose(1, 2).reshape(batch, length, width))) |
| return x + self.drop(self.mlp(self.norm2(x))) |
|
|
|
|
| class Eca(nn.Module): |
| """Decoder-only transformer. The output layer shares the token embedding.""" |
|
|
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
| self.tokens = nn.Embedding(config.vocab, config.width) |
| self.positions = nn.Embedding(config.block, config.width) |
| self.drop = nn.Dropout(config.dropout) |
| self.blocks = nn.Sequential(*(Block(config) for _ in range(config.layers))) |
| self.norm = nn.LayerNorm(config.width) |
| for embedding in (self.tokens, self.positions): |
| nn.init.normal_(embedding.weight, std=0.02) |
|
|
| def forward(self, ids): |
| """Next-token logits (batch, length, vocab).""" |
| x = self.tokens(ids) + self.positions(torch.arange(ids.shape[1], device=ids.device)) |
| return self.norm(self.blocks(self.drop(x))) @ self.tokens.weight.T |
|
|
|
|
| |
|
|
|
|
| def window_batch(data, args, generator): |
| """Random windows of the token stream and the same windows shifted by one.""" |
| block = args.block |
| start = torch.randint(len(data) - block - 1, (args.batch_size,), generator=generator) |
| x = torch.stack([data[i : i + block] for i in start]) |
| y = torch.stack([data[i + 1 : i + block + 1] for i in start]) |
| return x.to(DEVICE), y.to(DEVICE) |
|
|
|
|
| @torch.no_grad() |
| def validate(model, valid, batch_size=64): |
| """Mean loss per token and bits per character over non-overlapping windows.""" |
| valid, chars = valid |
| model.eval() |
| block = model.config.block |
| count = (len(valid) - 1) // block |
| x = valid[: count * block].view(count, block) |
| y = valid[1 : count * block + 1].view(count, block) |
| total = 0.0 |
| for xs, ys in zip(x.split(batch_size), y.split(batch_size)): |
| logits = model(xs.to(DEVICE)).float().reshape(-1, model.config.vocab) |
| total += F.cross_entropy(logits, ys.to(DEVICE).reshape(-1), reduction="sum").item() |
| model.train() |
| loss = total / (count * block) |
| return loss, loss * len(valid) / chars / math.log(2) |
|
|
|
|
| def update(model, optimizer, batch, lr): |
| """Take one optimizer step at the given learning rate and return the loss.""" |
| for group in optimizer.param_groups: |
| group["lr"] = lr |
| x, y = batch |
| with torch.autocast(DEVICE, dtype=torch.bfloat16, enabled=DEVICE == "cuda"): |
| logits = model(x) |
| loss = F.cross_entropy(logits.float().reshape(-1, model.config.vocab), y.reshape(-1)) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| optimizer.step() |
| return loss.item() |
|
|
|
|
| def schedule(args, step): |
| """Linear warmup then cosine decay to zero.""" |
| return args.lr * min(1, step / args.warmup) * 0.5 * (1 + math.cos(math.pi * step / args.steps)) |
|
|
|
|
| def train(args): |
| """Pre-train and keep the weights with the best validation score.""" |
| tokenizer, train_ids, valid = prepare(args) |
| torch.manual_seed(args.seed) |
| generator = torch.Generator().manual_seed(args.seed) |
| model = Eca(Config(args.vocab, args.block, args.width, args.layers, args.heads, args.dropout)) |
| model.to(DEVICE) |
| optimizer = torch.optim.AdamW(model.parameters(), args.lr, betas=(0.9, 0.95), weight_decay=0.1) |
| best = {"val_bpc": math.inf, "weights": copy.deepcopy(model.state_dict())} |
| start = time.time() |
| for step in range(1, args.steps + 1): |
| data = window_batch(train_ids, args, generator) |
| loss = update(model, optimizer, data, schedule(args, step)) |
| if step % args.eval_every == 0 or step == args.steps: |
| val_loss, val_bpc = validate(model, valid) |
| print(f"step {step} train {loss:.3f} val {val_loss:.3f} bpc {val_bpc:.3f}", flush=True) |
| if val_bpc < best["val_bpc"]: |
| best = { |
| "step": step, |
| "val_loss": round(val_loss, 4), |
| "val_bpc": round(val_bpc, 4), |
| "weights": copy.deepcopy(model.state_dict()), |
| } |
| model.load_state_dict(best.pop("weights")) |
| training = { |
| **best, |
| "params": sum(p.numel() for p in model.parameters()), |
| "steps": args.steps, |
| "batch_size": args.batch_size, |
| "lr": args.lr, |
| "train_tokens": len(train_ids), |
| "valid_tokens": len(valid[0]), |
| "valid_chars": valid[1], |
| "device": torch.cuda.get_device_name() if DEVICE == "cuda" else "cpu", |
| "hours": round((time.time() - start) / 3600, 2), |
| } |
| save(model, tokenizer, Path(args.output), training) |
| if args.push: |
| push(Path(args.output), args.repo, f"Train {args.steps} steps") |
|
|
|
|
| |
|
|
|
|
| def resolve(name): |
| """Local folder for a Hub repo id or a local path.""" |
| return Path(name) if Path(name).is_dir() else Path(snapshot_download(name)) |
|
|
|
|
| def load(name=REPO, device=DEVICE): |
| """Load a trained model and its tokenizer from a Hub repo id or a local folder.""" |
| folder = resolve(name) |
| config = Config(**json.loads((folder / "config.json").read_text(encoding="utf-8"))) |
| model = Eca(config) |
| model.load_state_dict(load_file(folder / "model.safetensors")) |
| return model.to(device).eval(), Tokenizer.from_file(str(folder / "tokenizer.json")) |
|
|
|
|
| @torch.no_grad() |
| def write(model, tokenizer, prompt, count=200, temperature=0.8): |
| """Continue the prompt by sampling `count` tokens.""" |
| ids = torch.tensor([tokenizer.encode(prompt).ids], device=next(model.parameters()).device) |
| for _ in range(count): |
| logits = model(ids[:, -model.config.block :])[:, -1] / temperature |
| logits = logits.masked_fill(logits < logits.topk(TOP_K).values[:, -1:], -torch.inf) |
| ids = torch.cat([ids, torch.multinomial(logits.softmax(-1), 1)], 1) |
| return tokenizer.decode(ids[0].tolist()) |
|
|
|
|
| def save(model, tokenizer, folder, training): |
| """Write weights, config, tokenizer, code and training results to a folder.""" |
| (folder / "results").mkdir(parents=True, exist_ok=True) |
| state = {key: value.cpu().contiguous() for key, value in model.state_dict().items()} |
| save_file(state, folder / "model.safetensors") |
| tokenizer.save(str(folder / "tokenizer.json")) |
| files = {"config.json": asdict(model.config), "results/train.json": training} |
| for name, value in files.items(): |
| (folder / name).write_text(json.dumps(value, indent=1) + "\n", encoding="utf-8") |
| if Path(__file__).resolve() != (folder / "eca.py").resolve(): |
| shutil.copy(__file__, folder / "eca.py") |
|
|
|
|
| def push(folder, repo, message): |
| """Upload a folder to a private model repo, creating it if needed.""" |
| api = HfApi() |
| api.create_repo(repo, private=True, exist_ok=True) |
| api.upload_folder(folder_path=folder, repo_id=repo, commit_message=message) |
|
|
|
|
| def card(args): |
| """Render card.jinja into card/README.md with the config and results stored in the repo.""" |
| files = {} |
| try: |
| folder = Path(snapshot_download(args.repo, allow_patterns=["config.json", "results/*"])) |
| files = { |
| path.relative_to(folder).as_posix(): json.loads(path.read_text(encoding="utf-8")) |
| for path in folder.glob("**/*.json") |
| } |
| except RepositoryNotFoundError: |
| pass |
| data = ModelCardData( |
| model_name=args.repo.split("/")[1], |
| language="pt", |
| license="mit", |
| library_name="pytorch", |
| pipeline_tag="text-generation", |
| datasets=[DATA], |
| tags=["portuguese", "literature", "from-scratch"], |
| ) |
| rendered = ModelCard.from_template( |
| data, |
| template_path=HERE / "card.jinja", |
| repo=args.repo, |
| data=DATA, |
| config=files.get("config.json"), |
| train=files.get("results/train.json"), |
| ) |
| (HERE / "card").mkdir(exist_ok=True) |
| rendered.save(HERE / "card" / "README.md") |
|
|
|
|
| def sample(args): |
| """Print a continuation of the prompt.""" |
| model, tokenizer = load(args.model) |
| print(write(model, tokenizer, args.prompt, args.tokens, args.temperature)) |
|
|
|
|
| def main(): |
| """Parse arguments and run a command.""" |
| parser = argparse.ArgumentParser(description=__doc__) |
| commands = parser.add_subparsers(dest="command", required=True) |
|
|
| train_parser = commands.add_parser("train") |
| train_parser.add_argument("--data", default=DATA, help="dataset repo id or a local text file") |
| train_parser.add_argument("--vocab", type=int, default=2048) |
| train_parser.add_argument("--block", type=int, default=256) |
| train_parser.add_argument("--width", type=int, default=256) |
| train_parser.add_argument("--layers", type=int, default=6) |
| train_parser.add_argument("--heads", type=int, default=4) |
| train_parser.add_argument("--dropout", type=float, default=0.1) |
| train_parser.add_argument("--steps", type=int, default=3000) |
| train_parser.add_argument("--batch-size", type=int, default=64) |
| train_parser.add_argument("--lr", type=float, default=1e-3) |
| train_parser.add_argument("--warmup", type=int, default=200) |
| train_parser.add_argument("--eval-every", type=int, default=250) |
| train_parser.add_argument("--seed", type=int, default=0) |
| train_parser.add_argument("--output", default="out/eca") |
| train_parser.add_argument("--repo", default=REPO) |
| train_parser.add_argument("--push", action="store_true") |
| train_parser.set_defaults(run=train) |
|
|
| write_parser = commands.add_parser("write") |
| write_parser.add_argument("prompt") |
| write_parser.add_argument("--model", default=REPO) |
| write_parser.add_argument("--tokens", type=int, default=200) |
| write_parser.add_argument("--temperature", type=float, default=0.8) |
| write_parser.set_defaults(run=sample) |
|
|
| card_parser = commands.add_parser("card") |
| card_parser.add_argument("--repo", default=REPO) |
| card_parser.set_defaults(run=card) |
|
|
| args = parser.parse_args() |
| args.run(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
