File size: 3,124 Bytes
e6dd5a3 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | # -*- coding: utf-8 -*-
"""
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()
|