| import torch | |
| import torch.nn as nn | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, embed_size, heads, ff_hidden_dim, dropout): | |
| super().__init__() | |
| self.attention = nn.MultiheadAttention(embed_dim=embed_size, num_heads=heads, batch_first=True) | |
| self.norm1 = nn.LayerNorm(embed_size) | |
| self.norm2 = nn.LayerNorm(embed_size) | |
| self.ff = nn.Sequential( | |
| nn.Linear(embed_size, ff_hidden_dim), | |
| nn.ReLU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(ff_hidden_dim, embed_size) | |
| ) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x): | |
| attn_output, _ = self.attention(x, x, x) | |
| x = self.norm1(x + self.dropout(attn_output)) | |
| ff_output = self.ff(x) | |
| x = self.norm2(x + self.dropout(ff_output)) | |
| return x | |
| class TransformerModel(nn.Module): | |
| def __init__(self, vocab_size, embed_size=512, num_heads=8, hidden_dim=2048, num_layers=6, max_len=512, dropout=0.1): | |
| super().__init__() | |
| self.embedding = nn.Embedding(vocab_size, embed_size) | |
| self.pos_embedding = nn.Parameter(torch.zeros(1, max_len, embed_size)) | |
| self.transformer_blocks = nn.Sequential( | |
| *[TransformerBlock(embed_size, num_heads, hidden_dim, dropout) for _ in range(num_layers)] | |
| ) | |
| self.norm = nn.LayerNorm(embed_size) | |
| self.output = nn.Linear(embed_size, vocab_size) | |
| def forward(self, x): | |
| seq_len = x.size(1) | |
| positions = self.pos_embedding[:, :seq_len, :] | |
| x = self.embedding(x) + positions | |
| x = self.transformer_blocks(x) | |
| x = self.norm(x) | |
| return self.output(x) | |