File size: 3,819 Bytes
e644e74 | 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 102 103 104 105 | import numpy as np
class TinyLM:
"""
Word-level language model.
Predicts the next word from the previous CONTEXT_SIZE words.
Supports batched forward/backward passes for fast training.
"""
def __init__(self, vocab_size, embed_dim=64, hidden_dim=128, context_size=5):
self.vocab_size = vocab_size
self.embed_dim = embed_dim
self.hidden_dim = hidden_dim
self.context_size = context_size
rng = np.random.default_rng(42)
self.embeddings = rng.standard_normal((vocab_size, embed_dim)).astype(np.float32) * 0.02
self.W1 = rng.standard_normal((embed_dim * context_size, hidden_dim)).astype(np.float32) * 0.02
self.b1 = np.zeros((1, hidden_dim), dtype=np.float32)
self.W2 = rng.standard_normal((hidden_dim, vocab_size)).astype(np.float32) * 0.02
self.b2 = np.zeros((1, vocab_size), dtype=np.float32)
def num_params(self):
return (self.embeddings.size + self.W1.size + self.b1.size
+ self.W2.size + self.b2.size)
def softmax(self, x):
exp = np.exp(x - np.max(x, axis=-1, keepdims=True))
return exp / np.sum(exp, axis=-1, keepdims=True)
def forward(self, token_ids):
"""Single-example forward pass (used at inference time)."""
embs = [self.embeddings[t] for t in token_ids]
x = np.concatenate(embs).reshape(1, -1)
h = np.tanh(np.dot(x, self.W1) + self.b1)
logits = np.dot(h, self.W2) + self.b2
probs = self.softmax(logits)
return x, h, probs
def forward_batch(self, token_id_batch):
"""
token_id_batch: (B, context_size) int array
Returns x (B, embed_dim*context_size), h (B, hidden_dim), probs (B, vocab_size)
"""
B = token_id_batch.shape[0]
x = self.embeddings[token_id_batch].reshape(B, -1)
h = np.tanh(x @ self.W1 + self.b1)
logits = h @ self.W2 + self.b2
probs = self.softmax(logits)
return x, h, probs
def train_step(self, token_ids, target, lr=0.01):
"""Single-example training step (kept for compatibility)."""
return self.train_batch(
np.array([token_ids], dtype=np.int64),
np.array([target], dtype=np.int64),
lr=lr,
)
def train_batch(self, token_id_batch, target_batch, lr=0.01):
"""
Vectorized batched training step.
token_id_batch: (B, context_size) int array
target_batch: (B,) int array
Returns mean loss over the batch.
"""
B = token_id_batch.shape[0]
x, h, probs = self.forward_batch(token_id_batch)
target_probs = probs[np.arange(B), target_batch]
loss = -np.mean(np.log(target_probs + 1e-9))
dout = probs.copy()
dout[np.arange(B), target_batch] -= 1.0
dout /= B
dW2 = h.T @ dout
db2 = dout.sum(axis=0, keepdims=True)
dh = dout @ self.W2.T
dtanh = (1.0 - h * h) * dh
dW1 = x.T @ dtanh
db1 = dtanh.sum(axis=0, keepdims=True)
dx = dtanh @ self.W1.T
self.W2 -= lr * dW2
self.b2 -= lr * db2
self.W1 -= lr * dW1
self.b1 -= lr * db1
dx = dx.reshape(B, self.context_size, self.embed_dim)
np.add.at(self.embeddings, token_id_batch, -lr * dx)
return loss
def sample(self, probs, temperature=0.6, top_k=10):
logits = np.log(np.clip(probs[0], 1e-9, 1.0)) / max(temperature, 0.05)
if top_k and top_k < len(logits):
idx = np.argpartition(logits, -top_k)[-top_k:]
mask = np.full_like(logits, -1e9)
mask[idx] = logits[idx]
logits = mask
exp = np.exp(logits - np.max(logits))
p = exp / np.sum(exp)
return int(np.random.choice(len(p), p=p))
|