File size: 3,772 Bytes
6b6e0dd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Train the tiny Qwen3-style model to generate Turkish names.
Run: python train.py
This script is intentionally linear and dependency-free (just torch). It reads
the cleaned names, builds a trivial character vocabulary inline, trains with a
plain loop, then samples a few names.
"""
import os
import torch
from config import ModelConfig
from model import TinyQwen
from tokenizer import CharTokenizer
# ---------------------------------------------------------------------------
# Hyperparameters
# ---------------------------------------------------------------------------
# Shared dataset lives one level up in ../data/ .
DATA_FILE = os.path.join(os.path.dirname(__file__), "..", "data", "koyler.txt")
BATCH_SIZE = 64
BLOCK_SIZE = 32 # context length used during training (<= cfg.max_seq_len)
STEPS = 8000
LEARNING_RATE = 3e-3
EVAL_EVERY = 200
SEED = 1337
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(SEED)
# ---------------------------------------------------------------------------
# Tokenizer (character level). See tokenizer.py.
# ---------------------------------------------------------------------------
tokenizer = CharTokenizer.from_file(DATA_FILE)
vocab_size = tokenizer.vocab_size
text = open(DATA_FILE, encoding="utf-8").read()
data = torch.tensor(tokenizer.encode(text), dtype=torch.long) # [N] whole corpus as ids
def get_batch():
"""Sample BATCH_SIZE random windows. Targets are inputs shifted by one."""
ix = torch.randint(len(data) - BLOCK_SIZE - 1, (BATCH_SIZE,))
x = torch.stack([data[i:i + BLOCK_SIZE] for i in ix])
y = torch.stack([data[i + 1:i + 1 + BLOCK_SIZE] for i in ix])
return x.to(device), y.to(device)
# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------
cfg = ModelConfig(vocab_size=vocab_size)
model = TinyQwen(cfg).to(device)
n_params = sum(p.numel() for p in model.parameters())
print(f"device={device} vocab_size={vocab_size} parameters={n_params:,}")
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
# ---------------------------------------------------------------------------
# Sampling helper: generate a few names starting from the newline token.
# ---------------------------------------------------------------------------
def sample_names(n: int = 10, max_new_tokens: int = 20):
model.eval()
start = torch.full((n, 1), tokenizer.newline_id, dtype=torch.long, device=device)
out = model.generate(start, max_new_tokens=max_new_tokens, temperature=1.0,
top_k=None, eos_id=tokenizer.eos_id)
model.train()
names = []
for row in out.tolist():
# Drop the leading newline, then keep up to the next newline.
s = tokenizer.decode(row[1:])
names.append(s.split("\n")[0])
return names
# ---------------------------------------------------------------------------
# Training loop
# ---------------------------------------------------------------------------
for step in range(1, STEPS + 1):
x, y = get_batch()
_, loss = model(x, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step % EVAL_EVERY == 0 or step == 1:
print(f"step {step:5d} loss {loss.item():.4f}")
print("\nbaseline loss (uniform guessing): %.4f" % (torch.log(torch.tensor(float(vocab_size)))))
print("\nsample names:")
for name in sample_names(10):
print(" ", name)
torch.save({"model": model.state_dict(), "chars": tokenizer.chars, "cfg": cfg},
"tiny_qwen_village.pt")
print("\nsaved checkpoint to tiny_qwen_village.pt")
|