File size: 3,256 Bytes
ce68c99 228f8cd ce68c99 228f8cd ce68c99 228f8cd | 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 | import math, torch
import torch.nn as nn, torch.nn.functional as F
T = 1024
N_LAYER = 24
N_HEAD = 12
N_EMBD = 768
VOCAB = 50304 # gpt2 vocab padded to a multiple of 64
class Block(nn.Module):
def __init__(s):
super().__init__()
s.ln1 = nn.LayerNorm(N_EMBD); s.ln2 = nn.LayerNorm(N_EMBD)
s.qkv = nn.Linear(N_EMBD, 3 * N_EMBD, bias=False)
s.proj = nn.Linear(N_EMBD, N_EMBD, bias=False)
s.fc = nn.Linear(N_EMBD, 4 * N_EMBD, bias=False)
s.out = nn.Linear(4 * N_EMBD, N_EMBD, bias=False)
def forward(s, x):
B, L, C = x.shape
q, k, v = s.qkv(s.ln1(x)).split(C, dim=2)
q, k, v = [t.view(B, L, N_HEAD, C // N_HEAD).transpose(1, 2) for t in (q, k, v)]
a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
x = x + s.proj(a.transpose(1, 2).reshape(B, L, C))
return x + s.out(F.gelu(s.fc(s.ln2(x))))
class GPT(nn.Module):
def __init__(s):
super().__init__()
s.wte = nn.Embedding(VOCAB, N_EMBD)
s.wpe = nn.Embedding(T, N_EMBD)
s.blocks = nn.ModuleList([Block() for _ in range(N_LAYER)])
s.lnf = nn.LayerNorm(N_EMBD)
s.head = nn.Linear(N_EMBD, VOCAB, bias=False)
s.head.weight = s.wte.weight # tied
s.apply(s._init)
for n, p in s.named_parameters():
if n.endswith('proj.weight') or n.endswith('out.weight'):
nn.init.normal_(p, std=0.02 / math.sqrt(2 * N_LAYER))
def _init(s, m):
if isinstance(m, (nn.Linear, nn.Embedding)):
nn.init.normal_(m.weight, std=0.02)
def forward(s, idx):
x = s.wte(idx) + s.wpe(torch.arange(idx.size(1), device=idx.device))
for b in s.blocks: x = b(x)
return s.head(s.lnf(x))
def load_iau1(path='model.safetensors', device='cpu'):
from safetensors.torch import load_file
m = GPT()
# head.weight is tied to wte.weight, so it is not stored in the file
m.load_state_dict(load_file(path), strict=False)
return m.to(device).eval()
@torch.no_grad()
def generate(model, prompt, n=80, temp=0.8, top_k=50, rep_penalty=1.2, no_repeat_ngram=3):
import tiktoken
enc = tiktoken.get_encoding('gpt2')
dev = next(model.parameters()).device
idx = torch.tensor([enc.encode_ordinary(prompt)], device=dev)
for _ in range(n):
logits = model(idx[:, -T:])[:, -1, :].float()
if rep_penalty != 1.0: # push down tokens that already appeared
seen = torch.unique(idx[0])
l = logits[0, seen]
logits[0, seen] = torch.where(l > 0, l / rep_penalty, l * rep_penalty)
if no_repeat_ngram and no_repeat_ngram >= 2 and idx.size(1) >= no_repeat_ngram:
k = no_repeat_ngram
toks = idx[0].tolist()
prefix = toks[-(k - 1):]
banned = [toks[i + k - 1] for i in range(len(toks) - k + 1) if toks[i:i + k - 1] == prefix]
if banned: logits[0, banned] = -float('inf') # never repeat an exact n-gram
logits = logits / temp
v, _ = torch.topk(logits, top_k)
logits[logits < v[:, [-1]]] = -float('inf')
idx = torch.cat([idx, torch.multinomial(F.softmax(logits, -1), 1)], dim=1)
return enc.decode(idx[0].tolist()) |