File size: 2,646 Bytes
267684e | 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 | #!/usr/bin/env python3
"""gary-4: a ~69KB chat model. Runs anywhere with just numpy.
Usage: python chat.py (interactive)
python chat.py "hi" (one-shot)"""
import json, sys, os
import numpy as np
D = os.path.dirname(os.path.abspath(__file__))
C = json.load(open(f"{D}/config.json"))
chars, E, H, L, BLK = C["chars"], C["n_embd"], C["n_head"], C["n_layer"], C["block_size"]
stoi = {c: i for i, c in enumerate(chars)}
z = np.load(f"{D}/gary4.int8.npz")
W = {k: z[k].astype(np.float32) * z[k + ".scale"] for k in z.files if not k.endswith(".scale")}
def ln(x, w, b, eps=1e-5):
m, v = x.mean(-1, keepdims=True), x.var(-1, keepdims=True)
return (x - m) / np.sqrt(v + eps) * w + b
def gelu(x):
return 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x**3)))
def forward(idx):
T = len(idx)
x = W["tok.weight"][idx] + W["pos.weight"][:T]
mask = np.triu(np.full((T, T), -1e9), 1)
hd = E // H
for i in range(L):
p = f"blocks.{i}."
h = ln(x, W[p+"ln1.weight"], W[p+"ln1.bias"])
qkv = h @ W[p+"attn.in_proj_weight"].T + W[p+"attn.in_proj_bias"]
q, k, v = np.split(qkv, 3, -1)
o = np.zeros_like(q)
for j in range(H):
s = slice(j*hd, (j+1)*hd)
att = q[:, s] @ k[:, s].T / np.sqrt(hd) + mask
att = np.exp(att - att.max(-1, keepdims=True))
att /= att.sum(-1, keepdims=True)
o[:, s] = att @ v[:, s]
x = x + o @ W[p+"attn.out_proj.weight"].T + W[p+"attn.out_proj.bias"]
h = ln(x, W[p+"ln2.weight"], W[p+"ln2.bias"])
h = gelu(h @ W[p+"mlp.0.weight"].T + W[p+"mlp.0.bias"])
x = x + h @ W[p+"mlp.2.weight"].T + W[p+"mlp.2.bias"]
x = ln(x, W["lnf.weight"], W["lnf.bias"])
return x[-1] @ W["tok.weight"].T
def generate(prompt, n=120, temp=0.7, seed=None):
rng = np.random.default_rng(seed)
idx = [stoi.get(c, stoi[" "]) for c in prompt]
out = ""
for _ in range(n):
logits = forward(idx[-BLK:]) / temp
p = np.exp(logits - logits.max()); p /= p.sum()
nxt = int(rng.choice(len(p), p=p))
ch = chars[nxt]
out += ch
idx.append(nxt)
if out.endswith("\n\n"): break
return out.strip()
def reply(msg):
return generate(f"User: {msg}\ngary-4:").split("\nUser:")[0].strip()
if __name__ == "__main__":
if len(sys.argv) > 1:
print(reply(" ".join(sys.argv[1:])))
else:
print("gary-4 (~69KB). ctrl-c to exit.")
while True:
try: msg = input("you: ")
except (EOFError, KeyboardInterrupt): break
print("gary-4:", reply(msg))
|