| 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 |
|
|
| 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 |
| 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() |
| |
| 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: |
| 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') |
| 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()) |
