From 836060ecdf731960aa9933027acb1e08ab865b7d Mon Sep 17 00:00:00 2001 From: Fendy Date: Thu, 16 Jul 2026 14:29:12 +0800 Subject: [PATCH] Initial commit: Simple Transformer implementation --- .gitignore | 50 ++++++++++++++++ README.md | 35 +++++++++++ example.py | 52 +++++++++++++++++ transformer.py | 156 +++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 293 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 example.py create mode 100644 transformer.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..28253cd --- /dev/null +++ b/.gitignore @@ -0,0 +1,50 @@ +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +*.manifest +*.spec +pip-log.txt +pip-delete-this-directory.txt +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ +*.log +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ +.idea/ +.vscode/ +*.swp +*.swo +*~ \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..9b3ff2e --- /dev/null +++ b/README.md @@ -0,0 +1,35 @@ +# Simple Transformer + +A minimal Transformer implementation using NumPy. + +## Files + +- `transformer.py` - Core Transformer model implementation +- `example.py` - Usage example + +## Usage + +```bash +python3 example.py +``` + +## Model Architecture + +The implementation includes: + +- **Multi-Head Attention**: Scaled dot-product attention with multiple heads +- **Feed-Forward Network**: Two-layer fully connected network with ReLU activation +- **Layer Normalization**: Applied after each sub-layer +- **Positional Encoding**: Sinusoidal position embeddings + +## Model Parameters + +Default configuration: +- Vocabulary size: 1000 +- Model dimension: 512 +- Number of heads: 8 +- Number of layers: 6 +- Feed-forward dimension: 2048 +- Max sequence length: 100 + +Total parameters: ~19.4M \ No newline at end of file diff --git a/example.py b/example.py new file mode 100644 index 0000000..148f479 --- /dev/null +++ b/example.py @@ -0,0 +1,52 @@ +import numpy as np +from transformer import SimpleTransformer, create_padding_mask, create_look_ahead_mask + + +def main(): + vocab_size = 1000 + d_model = 512 + num_heads = 8 + num_layers = 6 + d_ff = 2048 + max_seq_len = 100 + + model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len) + + batch_size = 2 + seq_len = 10 + x = np.random.randint(0, vocab_size, (batch_size, seq_len)) + + print("=== Simple Transformer Example ===") + print(f"Vocabulary size: {vocab_size}") + print(f"Model dimension: {d_model}") + print(f"Number of heads: {num_heads}") + print(f"Number of layers: {num_layers}") + print(f"Feed-forward dimension: {d_ff}") + print(f"Max sequence length: {max_seq_len}") + print() + + print(f"Input shape: {x.shape}") + print(f"Input sample: {x[0]}") + print() + + output = model.forward(x) + + print(f"Output shape: {output.shape}") + print(f"Output sample (first 5 values): {output[0, 0, :5]}") + print() + + total_params = model.count_parameters() + print(f"Total parameters: {total_params:,}") + print() + + print("=== Attention Mask Examples ===") + padding_mask = create_padding_mask(x) + print(f"Padding mask shape: {padding_mask.shape}") + + look_ahead_mask = create_look_ahead_mask(seq_len) + print(f"Look-ahead mask shape: {look_ahead_mask.shape}") + print(f"Look-ahead mask sample:\n{look_ahead_mask[:5, :5]}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/transformer.py b/transformer.py new file mode 100644 index 0000000..5c3794b --- /dev/null +++ b/transformer.py @@ -0,0 +1,156 @@ +import numpy as np + + +def softmax(x, axis=-1): + e_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) + return e_x / np.sum(e_x, axis=axis, keepdims=True) + + +def relu(x): + return np.maximum(0, x) + + +class Linear: + def __init__(self, in_features, out_features): + self.weight = np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features) + self.bias = np.zeros(out_features) + + def forward(self, x): + return x @ self.weight + self.bias + + +class LayerNorm: + def __init__(self, d_model, eps=1e-6): + self.gamma = np.ones(d_model) + self.beta = np.zeros(d_model) + self.eps = eps + + def forward(self, x): + mean = np.mean(x, axis=-1, keepdims=True) + std = np.std(x, axis=-1, keepdims=True) + return self.gamma * (x - mean) / (std + self.eps) + self.beta + + +class MultiHeadAttention: + def __init__(self, d_model, num_heads): + self.d_model = d_model + self.num_heads = num_heads + self.d_k = d_model // num_heads + + self.W_q = Linear(d_model, d_model) + self.W_k = Linear(d_model, d_model) + self.W_v = Linear(d_model, d_model) + self.W_o = Linear(d_model, d_model) + + def forward(self, Q, K, V, mask=None): + batch_size = Q.shape[0] + seq_len = Q.shape[1] + + Q = self.W_q.forward(Q).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) + K = self.W_k.forward(K).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) + V = self.W_v.forward(V).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3) + + attn_scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(self.d_k) + if mask is not None: + attn_scores = np.where(mask == 0, -1e9, attn_scores) + attn_probs = softmax(attn_scores, axis=-1) + attn_output = attn_probs @ V + + attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model) + output = self.W_o.forward(attn_output) + return output + + +class FeedForward: + def __init__(self, d_model, d_ff): + self.linear1 = Linear(d_model, d_ff) + self.linear2 = Linear(d_ff, d_model) + + def forward(self, x): + return self.linear2.forward(relu(self.linear1.forward(x))) + + +class TransformerBlock: + def __init__(self, d_model, num_heads, d_ff): + self.attention = MultiHeadAttention(d_model, num_heads) + self.feed_forward = FeedForward(d_model, d_ff) + self.norm1 = LayerNorm(d_model) + self.norm2 = LayerNorm(d_model) + + def forward(self, x, mask=None): + attn_output = self.attention.forward(x, x, x, mask) + x = self.norm1.forward(x + attn_output) + ff_output = self.feed_forward.forward(x) + x = self.norm2.forward(x + ff_output) + return x + + +class SimpleTransformer: + def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len): + self.d_model = d_model + self.embedding = np.random.randn(vocab_size, d_model) * 0.01 + self.positional_encoding = self._create_positional_encoding(max_seq_len, d_model) + self.transformer_blocks = [ + TransformerBlock(d_model, num_heads, d_ff) for _ in range(num_layers) + ] + + def _create_positional_encoding(self, max_seq_len, d_model): + pe = np.zeros((max_seq_len, d_model)) + position = np.arange(0, max_seq_len).reshape(-1, 1).astype(float) + div_term = np.exp(np.arange(0, d_model, 2).astype(float) * -(np.log(10000.0) / d_model)) + pe[:, 0::2] = np.sin(position * div_term) + pe[:, 1::2] = np.cos(position * div_term) + return pe + + def forward(self, x, mask=None): + seq_len = x.shape[1] + x = self.embedding[x] * np.sqrt(self.d_model) + x = x + self.positional_encoding[:seq_len, :] + + for block in self.transformer_blocks: + x = block.forward(x, mask) + + return x + + def count_parameters(self): + count = 0 + count += self.embedding.size + for block in self.transformer_blocks: + count += block.attention.W_q.weight.size + block.attention.W_q.bias.size + count += block.attention.W_k.weight.size + block.attention.W_k.bias.size + count += block.attention.W_v.weight.size + block.attention.W_v.bias.size + count += block.attention.W_o.weight.size + block.attention.W_o.bias.size + count += block.feed_forward.linear1.weight.size + block.feed_forward.linear1.bias.size + count += block.feed_forward.linear2.weight.size + block.feed_forward.linear2.bias.size + count += block.norm1.gamma.size + block.norm1.beta.size + count += block.norm2.gamma.size + block.norm2.beta.size + return count + + +def create_padding_mask(seq, pad_idx=0): + return (seq != pad_idx).astype(float)[:, np.newaxis, np.newaxis, :] + + +def create_look_ahead_mask(size): + mask = np.triu(np.ones((size, size)), k=1).astype(bool) + return (~mask).astype(float) + + +if __name__ == "__main__": + vocab_size = 1000 + d_model = 512 + num_heads = 8 + num_layers = 6 + d_ff = 2048 + max_seq_len = 100 + + model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len) + + batch_size = 2 + seq_len = 10 + x = np.random.randint(0, vocab_size, (batch_size, seq_len)) + + output = model.forward(x) + print(f"Input shape: {x.shape}") + print(f"Output shape: {output.shape}") + print(f"Model parameters: {model.count_parameters():,}") \ No newline at end of file