| import json, os |
| import torch, torch.nn as nn, torch.nn.functional as F |
| from tokenizers import Tokenizer |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, d, n_head): |
| super().__init__() |
| self.n_head = n_head |
| self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d) |
| self.qkv = nn.Linear(d, 3 * d, bias=False) |
| self.proj = nn.Linear(d, d, bias=False) |
| self.fc = nn.Linear(d, 4 * d, bias=False) |
| self.fc2 = nn.Linear(4 * d, d, bias=False) |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
| q, k, v = self.qkv(self.ln1(x)).split(C, dim=2) |
| q, k, v = [t.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) for t in (q, k, v)] |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| x = x + self.proj(y.transpose(1, 2).reshape(B, T, C)) |
| return x + self.fc2(F.gelu(self.fc(self.ln2(x)))) |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| d = cfg["n_embd"] |
| self.ctx = cfg["ctx"] |
| self.tok_emb = nn.Embedding(cfg["vocab"], d) |
| self.pos_emb = nn.Embedding(cfg["ctx"], d) |
| self.blocks = nn.ModuleList([Block(d, cfg["n_head"]) for _ in range(cfg["n_layer"])]) |
| self.ln_f = nn.LayerNorm(d) |
| self.head = nn.Linear(d, cfg["vocab"], bias=False) |
| self.head.weight = self.tok_emb.weight |
|
|
| def forward(self, idx): |
| x = self.tok_emb(idx) + self.pos_emb(torch.arange(idx.shape[1], device=idx.device)) |
| for blk in self.blocks: |
| x = blk(x) |
| return self.head(self.ln_f(x)) |
|
|
| @torch.no_grad() |
| def generate(self, idx, max_new=100, temperature=0.7, top_k=40, stop_id=None): |
| for _ in range(max_new): |
| logits = self(idx[:, -self.ctx:])[:, -1, :] / max(temperature, 1e-5) |
| if top_k: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = -float("inf") |
| nxt = torch.multinomial(F.softmax(logits, dim=-1), 1) |
| idx = torch.cat([idx, nxt], dim=1) |
| if stop_id is not None and nxt.item() == stop_id: |
| break |
| return idx |
|
|
|
|
| def load(path, device="cpu"): |
| with open(os.path.join(path, "config.json")) as f: |
| cfg = json.load(f) |
| model = GPT(cfg) |
| model.load_state_dict(torch.load(os.path.join(path, "model.pt"), map_location=device)) |
| tok = Tokenizer.from_file(os.path.join(path, "tokenizer.json")) |
| return model.to(device).eval(), tok |
|
|
|
|
| def reply(model, tok, history, user_msg, temperature=0.7, top_k=40, max_new=120): |
| |
| USER, ASSIST, END = [tok.token_to_id(t) for t in ("<|user|>", "<|assistant|>", "<|end|>")] |
| ids = [] |
| for u, a in history: |
| ids += [USER] + tok.encode(u).ids + [END, ASSIST] + tok.encode(a).ids + [END] |
| ids += [USER] + tok.encode(user_msg).ids + [END, ASSIST] |
| ids = ids[-(model.ctx - 1):] |
| device = next(model.parameters()).device |
| out = model.generate(torch.tensor([ids], device=device), max_new=max_new, |
| temperature=temperature, top_k=top_k, stop_id=END) |
| new = out[0, len(ids):].tolist() |
| if new and new[-1] == END: |
| new = new[:-1] |
| return tok.decode(new).strip() |
|
|
