Spaces:
Running
Running
| # bot_model.py | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, d_model, num_heads, dropout=0.1): | |
| super().__init__() | |
| assert d_model % num_heads == 0 | |
| self.num_heads = num_heads | |
| self.d_k = d_model // num_heads | |
| # NOTE: match checkpoint names exactly | |
| self.W_q = nn.Linear(d_model, d_model) | |
| self.W_k = nn.Linear(d_model, d_model) | |
| self.W_v = nn.Linear(d_model, d_model) | |
| self.W_o = nn.Linear(d_model, d_model) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x): | |
| B, T, C = x.size() | |
| q = self.W_q(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2) | |
| k = self.W_k(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2) | |
| v = self.W_v(x).view(B, T, self.num_heads, C//self.num_heads).transpose(1, 2) | |
| att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k) | |
| mask = torch.tril(torch.ones(T, T, device=x.device)) | |
| att = att.masked_fill(mask == 0, float('-inf')) | |
| att = F.softmax(att, dim=-1) | |
| att = self.dropout(att) | |
| y = att @ v | |
| y = y.transpose(1, 2).contiguous().view(B, T, C) | |
| return self.W_o(y) | |
| class GPTBlock(nn.Module): | |
| def __init__(self, d_model, num_heads, d_ff=768, dropout=0.1): | |
| super().__init__() | |
| self.ln1 = nn.LayerNorm(d_model) | |
| self.attn = CausalSelfAttention(d_model, num_heads, dropout) | |
| self.ln2 = nn.LayerNorm(d_model) | |
| # NOTE: match checkpoint module name "ffn" | |
| self.ffn = nn.Sequential( | |
| nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), nn.Dropout(dropout) | |
| ) | |
| def forward(self, x): | |
| x = x + self.attn(self.ln1(x)) | |
| x = x + self.ffn(self.ln2(x)) | |
| return x | |
| class GPT(nn.Module): | |
| def __init__(self, vocab_size, d_model=192, num_heads=6, num_layers=4, max_len=320, dropout=0.1): | |
| super().__init__() | |
| self.token_emb = nn.Embedding(vocab_size, d_model) | |
| self.pos_emb = nn.Embedding(max_len, d_model) | |
| self.blocks = nn.ModuleList([GPTBlock(d_model, num_heads, 4*d_model, dropout) for _ in range(num_layers)]) | |
| self.ln_f = nn.LayerNorm(d_model) | |
| self.head = nn.Linear(d_model, vocab_size, bias=False) | |
| self.cfg = dict(vocab_size=vocab_size, d_model=d_model, num_heads=num_heads, num_layers=num_layers, max_len=max_len, dropout=dropout) | |
| def forward(self, idx): | |
| B, T = idx.shape | |
| pos = torch.arange(0, T, device=idx.device).unsqueeze(0) | |
| x = self.token_emb(idx) + self.pos_emb(pos) | |
| for blk in self.blocks: | |
| x = blk(x) | |
| x = self.ln_f(x) | |
| return self.head(x) | |