| |
| """ |
| Z-Born 4.5M — живой чат на CPU. |
| |
| pip install torch tokenizers |
| python demo.py "Once upon a time" |
| """ |
| import os |
| import sys |
|
|
| os.environ['CUDA_VISIBLE_DEVICES'] = '' |
|
|
| if hasattr(sys.stdout, 'reconfigure'): |
| sys.stdout.reconfigure(encoding='utf-8', errors='replace') |
|
|
| import torch |
| import torch.nn as nn |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
|
|
|
|
| class BornFactorRNN(nn.Module): |
| """1L tanh-RNN, рождённый в Z-форме: Wx/Wh ранга k (U,S,V).""" |
|
|
| def __init__(self, v, hid, k): |
| super().__init__() |
| self.v, self.hid, self.k = v, hid, k |
| self.Ux = nn.Parameter(torch.zeros(hid, k)) |
| self.Vx = nn.Parameter(torch.zeros(v, k)) |
| self.Sx = nn.Parameter(torch.zeros(k)) |
| self.Uh = nn.Parameter(torch.zeros(hid, k)) |
| self.Vh = nn.Parameter(torch.zeros(hid, k)) |
| self.Sh = nn.Parameter(torch.zeros(k)) |
| self.b_ih = nn.Parameter(torch.zeros(hid)) |
| self.b_hh = nn.Parameter(torch.zeros(hid)) |
| self.head = nn.Linear(hid, v) |
|
|
| def hidden(self, xt, h): |
| xin = (self.Vx[xt] * self.Sx) @ self.Ux.T |
| hh = ((h @ self.Vh) * self.Sh) @ self.Uh.T |
| return torch.tanh(xin + hh + self.b_ih + self.b_hh) |
|
|
| def step(self, xt, h): |
| h = self.hidden(xt, h) |
| return self.head(h), h |
|
|
| @torch.no_grad() |
| def generate(self, ids, n, temp=0.8, top_k=40, rep=1.12): |
| self.eval() |
| h = torch.zeros(1, self.hid) |
| logits = None |
| for tok in ids: |
| logits, h = self.step(torch.tensor([tok]), h) |
| out = list(ids) |
| for _ in range(n): |
| lg = logits[0] |
| if rep and out: |
| lg = lg.clone() |
| for t in set(out[-48:]): |
| lg[t] = lg[t] / rep if lg[t] > 0 else lg[t] * rep |
| lg = lg / max(temp, 1e-6) |
| if top_k and top_k < lg.numel(): |
| thr = torch.topk(lg, top_k).values[-1] |
| lg = lg.masked_fill(lg < thr, float('-inf')) |
| nxt = torch.multinomial(torch.softmax(lg, -1), 1) |
| out.append(int(nxt)) |
| logits, h = self.step(nxt.view(1), h) |
| return out[len(ids):] |
|
|
|
|
| def load(): |
| ckpt = os.path.join(HERE, 'zborn.pt') |
| if not os.path.isfile(ckpt): |
| raise SystemExit('нет %s' % ckpt) |
| pack = torch.load(ckpt, map_location='cpu', weights_only=False) |
| meta = pack['meta'] |
| from tokenizers import Tokenizer |
| tok = Tokenizer.from_file(os.path.join(HERE, 'tokenizer.json')) |
| m = BornFactorRNN(meta['v'], meta['hid'], meta['k']) |
| m.load_state_dict(pack['model']) |
| m.eval() |
| return m, tok, meta |
|
|
|
|
| def main(): |
| prompt = ' '.join(sys.argv[1:]).strip() or 'The meaning of life is' |
| m, tok, meta = load() |
| print('Z-Born %s hid=%s k=%s n=%s step=%s CPU' |
| % (meta.get('kind', '?'), meta.get('hid'), meta.get('k'), |
| meta.get('n_born'), meta.get('step'))) |
| ids = tok.encode(prompt).ids |
| out = m.generate(ids, n=120) |
| print('PROMPT:', prompt) |
| print(tok.decode(out)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
