Initial commit: Simple Transformer implementation

2026-07-16 14:29:12 +08:00
commit 836060ecdf
4 changed files with 293 additions and 0 deletions
@@ -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
+*~
@@ -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
@@ -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()
@@ -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():,}")