Spaces:
Running
Running
File size: 2,778 Bytes
dd2fdf7 71807fe dd2fdf7 71807fe dd2fdf7 71807fe dd2fdf7 71807fe dd2fdf7 71807fe dd2fdf7 71807fe dd2fdf7 71807fe dd2fdf7 71807fe dd2fdf7 71807fe dd2fdf7 71807fe | 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 | # 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)
|