File size: 3,837 Bytes
591ec58 |
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 |
# model.py
import torch
import torch.nn as nn
import math
class MultiHeadSelfAttention(nn.Module):
def __init__(self, hidden_size, num_heads, dropout=0.1):
super().__init__()
assert hidden_size % num_heads == 0, "hidden_size must be divisible by num_heads"
self.hidden_size = hidden_size
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.query = nn.Linear(hidden_size, hidden_size)
self.key = nn.Linear(hidden_size, hidden_size)
self.value = nn.Linear(hidden_size, hidden_size)
self.out = nn.Linear(hidden_size, hidden_size)
self.dropout = nn.Dropout(dropout)
self.scale = math.sqrt(self.head_dim)
def forward(self, x, mask=None, padding_mask=None):
batch_size, seq_len, _ = x.size()
q = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
k = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
v = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale
if mask is not None:
scores = scores.masked_fill(mask == 1, -1e4) # Adjusted for FP16 compatibility
if padding_mask is not None:
padding_mask = padding_mask.unsqueeze(1).unsqueeze(2)
scores = scores.masked_fill(padding_mask, -1e4) # Adjusted for FP16 compatibility
attn = torch.softmax(scores, dim=-1)
attn = self.dropout(attn)
out = torch.matmul(attn, v).transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_size)
out = self.out(out)
return out
class TransformerLayer(nn.Module):
def __init__(self, hidden_size, num_heads, dropout=0.1):
super().__init__()
self.attn = MultiHeadSelfAttention(hidden_size, num_heads, dropout)
self.ffn = nn.Sequential(
nn.Linear(hidden_size, 4 * hidden_size),
nn.ReLU(),
nn.Linear(4 * hidden_size, hidden_size),
nn.Dropout(dropout)
)
self.ln1 = nn.LayerNorm(hidden_size)
self.ln2 = nn.LayerNorm(hidden_size)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None, padding_mask=None):
x = self.ln1(x)
attn_out = self.attn(x, mask, padding_mask)
x = x + self.dropout(attn_out)
x = self.ln2(x)
ffn_out = self.ffn(x)
x = x + self.dropout(ffn_out)
return x
class TransformerModel(nn.Module):
def __init__(self, vocab_size, hidden_size=512, num_layers=6, num_heads=8, dropout=0.1):
super().__init__()
self.token_embedding = nn.Embedding(vocab_size, hidden_size)
self.pos_embedding = nn.Embedding(512, hidden_size) # Fixed max_seq_len=512
self.layers = nn.ModuleList([
TransformerLayer(hidden_size, num_heads, dropout) for _ in range(num_layers)
])
self.final_ln = nn.LayerNorm(hidden_size)
self.head = nn.Linear(hidden_size, vocab_size)
self.dropout = nn.Dropout(dropout)
def forward(self, input_ids, padding_mask=None):
batch_size, seq_len = input_ids.size()
positions = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0).expand_as(input_ids)
x = self.token_embedding(input_ids) + self.pos_embedding(positions)
x = self.dropout(x)
causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=input_ids.device), diagonal=1).bool()
for layer in self.layers:
x = layer(x, causal_mask, padding_mask)
x = self.final_ln(x)
logits = self.head(x)
return logits |