Upload 6 files

train/config.yaml

@@ -0,0 +1,19 @@
+vocab_size: 16000
+block_size: 256
+n_layer: 6
+n_head: 6
+n_embed: 384
+batch_size: 32
+micro_batches: 4
+lr: 3.0e-4
+min_lr: 3.0e-5
+warmup_steps: 200
+max_steps: 1000
+weight_decay: 0.01
+grad_clip: 1.0
+dtype: "float32"
+device: "auto"
+save_dir: "out/pretrain"
+tokenizer_path: "out/tokenizer.json"
+train_txt: "data/corpus_raw.txt"
+sft_jsonl: "data/sft_train.jsonl"

train/data_utils.py

@@ -0,0 +1,13 @@
+import torch
+from torch.utils.data import Dataset
+
+class TextDataset(Dataset):
+    def __init__(self, ids, block_size):
+        self.ids = ids
+        self.block = block_size
+    def __len__(self):
+        return max(1, len(self.ids) - self.block)
+    def __getitem__(self, i):
+        x = self.ids[i:i+self.block]
+        y = self.ids[i+1:i+self.block+1]
+        return torch.tensor(x, dtype=torch.long), torch.tensor(y, dtype=torch.long)

train/gen_sample.py

@@ -0,0 +1,23 @@
+import torch, argparse, json
+from tokenizers import Tokenizer
+from model.tiny_gpt2 import TinyGPT2, GPTConfig
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--prompt", type=str, required=True)
+parser.add_argument("--ckpt", type=str, default="out/sft/model_sft.pt")
+parser.add_argument("--cfg", type=str, default="out/pretrain/gpt_config.json")
+parser.add_argument("--tok", type=str, default="out/tokenizer.json")
+args = parser.parse_args()
+
+tok = Tokenizer.from_file(args.tok)
+cfg = GPTConfig(**json.load(open(args.cfg)))
+m = TinyGPT2(cfg)
+m.load_state_dict(torch.load(args.ckpt, map_location="cpu"))
+m.eval()
+
+ids = tok.encode("[BOS] " + args.prompt).ids
+x = torch.tensor([ids], dtype=torch.long)
+with torch.no_grad():
+    y = m.generate(x, max_new_tokens=80)
+text = tok.decode(y[0].tolist())
+print(text)

train/prepare_corpus.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+
+SRC = Path("data/corpus_raw.txt")
+if __name__ == "__main__":
+    text = SRC.read_text(encoding="utf-8", errors="ignore")
+    text = text.replace("\r\n", "\n").strip()
+    SRC.write_text(text, encoding="utf-8")
+    print("cleaned corpus in-place.")

train/pretrain.py

@@ -0,0 +1,64 @@
+import yaml, math, time, json
+import torch
+from pathlib import Path
+from tokenizers import Tokenizer
+from torch.utils.data import DataLoader
+from torch.optim import AdamW
+from model.tiny_gpt2 import TinyGPT2, GPTConfig
+from train.data_utils import TextDataset
+
+def get_device(name):
+    if name == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    return name
+
+def cosine_lr(step, max_steps, base, min_lr, warmup):
+    if step < warmup:
+        return base * step / max(1, warmup)
+    progress = (step - warmup)/max(1, max_steps - warmup)
+    return min_lr + 0.5*(base-min_lr)*(1+math.cos(math.pi*progress))
+
+if __name__ == "__main__":
+    cfg = yaml.safe_load(open("train/config.yaml"))
+    device = get_device(cfg["device"])
+    Path(cfg["save_dir"]).mkdir(parents=True, exist_ok=True)
+
+    tok = Tokenizer.from_file(cfg["tokenizer_path"])
+    ids = tok.encode(open(cfg["train_txt"], "r", encoding="utf-8").read()).ids
+    ds = TextDataset(ids, cfg["block_size"])
+    dl = DataLoader(ds, batch_size=cfg["batch_size"], shuffle=True, drop_last=True)
+
+    gcfg = GPTConfig(
+        vocab_size=cfg["vocab_size"],
+        n_layer=cfg["n_layer"],
+        n_head=cfg["n_head"],
+        n_embed=cfg["n_embed"],
+        block_size=cfg["block_size"],
+    )
+    model = TinyGPT2(gcfg).to(device)
+
+    opt = AdamW(model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"])
+    step, t0 = 0, time.time()
+    model.train()
+    for epoch in range(999999):
+        for x, y in dl:
+            step += 1
+            x, y = x.to(device), y.to(device)
+            logits = model(x)
+            loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["grad_clip"])
+            lr = cosine_lr(step, cfg["max_steps"], cfg["lr"], cfg["min_lr"], cfg["warmup_steps"])
+            for g in opt.param_groups: g["lr"] = lr
+            opt.step(); opt.zero_grad(set_to_none=True)
+
+            if step % 100 == 0:
+                dt = time.time() - t0; t0 = time.time()
+                print(f"step {step:6d} | loss {loss.item():.4f} | lr {lr:.2e} | {dt:.2f}s")
+
+            if step >= cfg["max_steps"]:
+                torch.save(model.state_dict(), f"{cfg['save_dir']}/model.pt")
+                with open(f"{cfg['save_dir']}/gpt_config.json", "w") as f:
+                    json.dump(gcfg.__dict__, f, indent=2)
+                print("saved checkpoint. done.")
+                raise SystemExit

train/sft.py

@@ -0,0 +1,52 @@
+import json, yaml, time
+import torch
+from pathlib import Path
+from tokenizers import Tokenizer
+from torch.utils.data import Dataset, DataLoader
+from torch.optim import AdamW
+from model.tiny_gpt2 import TinyGPT2, GPTConfig
+
+class SFTDataset(Dataset):
+    def __init__(self, jsonl_path, tokenizer, block_size):
+        self.block = block_size
+        self.tok = tokenizer
+        self.samples = [json.loads(l) for l in open(jsonl_path, 'r', encoding='utf-8')]
+        self.ids = []
+        for s in self.samples:
+            text = f"Instruction:\n{s['instruction'].strip()}\nAnswer:\n{s['output'].strip()}\n"
+            self.ids.append(self.tok.encode(text).ids)
+    def __len__(self): return len(self.ids)
+    def __getitem__(self, i):
+        ids = self.ids[i][:self.block]
+        x = ids[:-1]; y = ids[1:]
+        return torch.tensor(x, dtype=torch.long), torch.tensor(y, dtype=torch.long)
+
+if __name__ == "__main__":
+    cfg = yaml.safe_load(open("train/config.yaml"))
+    Path("out/sft").mkdir(parents=True, exist_ok=True)
+    tok = Tokenizer.from_file(cfg["tokenizer_path"])
+
+    gcfg = GPTConfig(**json.load(open(Path(cfg["save_dir"]) / "gpt_config.json")))
+    model = TinyGPT2(gcfg)
+    model.load_state_dict(torch.load(Path(cfg["save_dir"])/"model.pt", map_location="cpu"))
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+
+    ds = SFTDataset(cfg["sft_jsonl"], tok, gcfg.block_size)
+    dl = DataLoader(ds, batch_size=8, shuffle=True, drop_last=True)
+    opt = AdamW(model.parameters(), lr=1e-4)
+
+    model.train()
+    t0 = time.time()
+    for step, (x,y) in enumerate(dl, start=1):
+        x,y = x.to(device), y.to(device)
+        logits = model(x)
+        loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
+        loss.backward(); opt.step(); opt.zero_grad(set_to_none=True)
+        if step % 50 == 0:
+            dt = time.time()-t0; t0=time.time()
+            print(f"sft step {step:5d} | loss {loss.item():.4f} | {dt:.2f}s")
+        if step >= 800: break
+
+    torch.save(model.state_dict(), "out/sft/model_sft.pt")
+    print("SFT saved.")