| import torch
|
| import torch.nn.functional as F
|
| import torch.nn as nn
|
| import math
|
|
|
| embedding = 256
|
| heads = 4
|
| layers = 4
|
| dropout = 0.1
|
| msl = 160
|
|
|
| class PositionalEncoding(nn.Module):
|
|
|
| def __init__(self, d_model, max_len=5000):
|
|
|
| super(PositionalEncoding, self).__init__()
|
|
|
| pe = torch.zeros(max_len, d_model)
|
|
|
| position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
|
|
|
| div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
|
|
|
| pe[:, 0::2] = torch.sin(position * div_term)
|
|
|
| pe[:, 1::2] = torch.cos(position * div_term)
|
|
|
| pe = pe.unsqueeze(0).transpose(0, 1)
|
|
|
| self.register_buffer('pe', pe)
|
|
|
|
|
| def forward(self, x):
|
|
|
| return x + self.pe[:x.size(1), :].transpose(0, 1)
|
|
|
|
|
| class CausalSelfAttention(nn.Module):
|
| def __init__(self, d_model, nhead, dropout=0.1):
|
| super().__init__()
|
|
|
| assert d_model % nhead == 0
|
|
|
| self.nhead = nhead
|
| self.head_dim = d_model // nhead
|
| self.dropout = dropout
|
|
|
| self.qkv = nn.Linear(d_model, d_model * 3)
|
| self.out_proj = nn.Linear(d_model, d_model)
|
|
|
| def forward(self, x):
|
| B, T, C = x.shape
|
|
|
|
|
| q, k, v = self.qkv(x).chunk(3, dim=-1)
|
|
|
|
|
| q = q.view(B, T, self.nhead, self.head_dim).transpose(1, 2)
|
| k = k.view(B, T, self.nhead, self.head_dim).transpose(1, 2)
|
| v = v.view(B, T, self.nhead, self.head_dim).transpose(1, 2)
|
|
|
| y = F.scaled_dot_product_attention(
|
| q,
|
| k,
|
| v,
|
| attn_mask=None,
|
| dropout_p=self.dropout if self.training else 0.0,
|
| is_causal=True
|
| )
|
|
|
| y = y.transpose(1, 2).contiguous().view(B, T, C)
|
|
|
| return self.out_proj(y)
|
|
|
|
|
| class TransformerBlock(nn.Module):
|
| def __init__(self, d_model, nhead, dropout=0.1):
|
| super().__init__()
|
|
|
| self.norm1 = nn.LayerNorm(d_model)
|
| self.attention = CausalSelfAttention(
|
| d_model,
|
| nhead,
|
| dropout
|
| )
|
|
|
| self.norm2 = nn.LayerNorm(d_model)
|
|
|
| self.ffn = nn.Sequential(
|
| nn.Linear(d_model, d_model * 4),
|
| nn.GELU(),
|
| nn.Linear(d_model * 4, d_model),
|
| nn.Dropout(dropout)
|
| )
|
|
|
| def forward(self, x):
|
| x = x + self.attention(self.norm1(x))
|
|
|
| x = x + self.ffn(self.norm2(x))
|
|
|
| return x
|
|
|
|
|
| class TransformerLanguageModel(nn.Module):
|
| def __init__(
|
| self,
|
| vocab_size,
|
| d_model=512,
|
| nhead=8,
|
| num_layers=8,
|
| dropout=0.1,
|
| max_seq_len=160
|
| ):
|
| super().__init__()
|
|
|
| self.d_model = d_model
|
| self.max_seq_len = max_seq_len
|
|
|
| self.token_embedding = nn.Embedding(
|
| vocab_size,
|
| d_model
|
| )
|
|
|
| self.positional_encoding = PositionalEncoding(
|
| d_model,
|
| max_seq_len
|
| )
|
|
|
| self.transformer = nn.ModuleList([
|
| TransformerBlock(
|
| d_model,
|
| nhead,
|
| dropout
|
| )
|
| for _ in range(num_layers)
|
| ])
|
|
|
| self.final_norm = nn.LayerNorm(d_model)
|
|
|
| self.output_layer = nn.Linear(
|
| d_model,
|
| vocab_size,
|
| bias=False
|
| )
|
|
|
| def forward(self, src):
|
| x = self.token_embedding(src)
|
|
|
| x = self.positional_encoding(x)
|
|
|
| for layer in self.transformer:
|
| x = layer(x)
|
|
|
| x = self.final_norm(x)
|
|
|
| return self.output_layer(x) |
