model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class Attention(nn.Module):
    def __init__(self, dim, head_dim, dropout):
        super().__init__()
        self.head_dim = head_dim
        self.w_q = nn.Linear(dim, head_dim)
        self.w_k = nn.Linear(dim, head_dim)
        self.w_v = nn.Linear(dim, head_dim)
        self.attn_dropout = nn.Dropout(dropout)

    def forward(self, x):
        q = self.w_q(x)
        k = self.w_k(x)
        v = self.w_v(x)
        attn_scores = q @ k.transpose(-1, -2) / math.sqrt(self.head_dim)
        mask = torch.triu(torch.full((x.shape[1], x.shape[1]), float('-inf'), device=x.device), diagonal=1)
        attn_scores = attn_scores + mask
        attn_w = self.attn_dropout(F.softmax(attn_scores, dim=-1))
        return attn_w @ v


class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, dim, dropout):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.heads = nn.ModuleList([Attention(dim, self.head_dim, dropout) for _ in range(num_heads)])
        self.multihead_proj = nn.Linear(dim, dim)
        self.proj_dropout = nn.Dropout(dropout)

    def forward(self, x):
        head_outs = [head(x) for head in self.heads]
        multihead_out = torch.cat(head_outs, dim=-1)
        return self.proj_dropout(self.multihead_proj(multihead_out))


class FFN(nn.Module):
    def __init__(self, dim, expansion, dropout):
        super().__init__()
        self.fc1 = nn.Linear(dim, dim * expansion)
        self.fc2 = nn.Linear(dim * expansion, dim)
        self.gelu = nn.GELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.dropout(self.fc2(self.gelu(self.fc1(x))))


class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, expansion, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(dim)
        self.multihead_attn = MultiHeadAttention(num_heads, dim, dropout)
        self.ln2 = nn.LayerNorm(dim)
        self.ffn = FFN(dim, expansion, dropout)

    def forward(self, x):
        x = x + self.multihead_attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self, vocab_size, dim, num_heads, expansion, num_layers, context_length, dropout=0.1):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, dim)
        self.pos_embedding = nn.Embedding(context_length, dim)
        self.emb_dropout = nn.Dropout(dropout)
        self.transformer_blocks = nn.ModuleList([
            TransformerBlock(dim, num_heads, expansion, dropout) for _ in range(num_layers)
        ])
        self.ln_head = nn.LayerNorm(dim)
        self.linear_head = nn.Linear(dim, vocab_size)

    def forward(self, x):
        tok_emb = self.token_embedding(x)
        pos_emb = self.pos_embedding(torch.arange(x.shape[1], device=x.device))
        x = self.emb_dropout(tok_emb + pos_emb)
        for block in self.transformer_blocks:
            x = block(x)
        x = self.ln_head(x)
        return self.linear_head(x)