Release TONER-1 Pro: our own 12M model and tokenizer as a single PostScript print job
20aa772 verified | # Export llama2.c stories260K checkpoint to int8 (per-row) weights embedded in a PostScript file. | |
| import struct, sys, array | |
| ckpt, tokf, out = sys.argv[1], sys.argv[2], sys.argv[3] | |
| raw = open(ckpt,'rb').read() | |
| dim, hid, L, nh, nkv, V, seq = struct.unpack('7i', raw[:28]) | |
| shared = V > 0; V = abs(V) | |
| f = array.array('f'); f.frombytes(raw[28:]) | |
| hs = dim//nh; kvd = dim*nkv//nh | |
| p = 0 | |
| def take(n): | |
| global p | |
| a = f[p:p+n]; p += n; return a | |
| emb = take(V*dim); rms_att = take(L*dim); wq = take(L*dim*dim); wk = take(L*dim*kvd); wv = take(L*dim*kvd) | |
| wo = take(L*dim*dim); rms_ffn = take(L*dim); w1 = take(L*hid*dim); w2 = take(L*dim*hid); w3 = take(L*hid*dim) | |
| rms_f = take(dim) | |
| assert shared | |
| nparams = V*dim + L*(2*dim*dim + 2*dim*kvd + 3*dim*hid + 2*dim) + dim | |
| o = open(out,'w') | |
| def psarr(items, chunk=200): | |
| # build arrays in chunks so the operand stack never exceeds ~200 items (Adobe limit is 500) | |
| if len(items) <= chunk: return '[' + ' '.join(items) + ']' | |
| out = '%d array' % len(items) | |
| for i in range(0, len(items), chunk): | |
| out += ' dup %d [%s] putinterval\n' % (i, ' '.join(items[i:i+chunk])) | |
| return out | |
| def realarr(a): return psarr(['%.6g' % v for v in a]) | |
| def qrows(a, rows, cols): | |
| # returns PS text: [ [scales] [<hex> ...] ] | |
| sc=[]; hx=[] | |
| for r in range(rows): | |
| row = a[r*cols:(r+1)*cols] | |
| m = max(abs(v) for v in row) or 1e-8 | |
| s = m/127.0 | |
| b = bytes(max(1,min(255,int(round(v/s))+128)) for v in row) | |
| sc.append(s); hx.append('<'+b.hex()+'>') | |
| return '[' + realarr(sc) + '\n' + psarr(hx) + ']' | |
| o.write('%!PS-Adobe-3.0\n%%Title: TONER-1 (stories260K int8, PostScript inference)\n') | |
| o.write('/D %d def /HID %d def /NL %d def /NH %d def /NKV %d def /V %d def /SEQ %d def /HS %d def /KVD %d def /KVMUL %d def /NPARAMS %d def\n' % (dim,hid,L,nh,nkv,V,seq,hs,kvd,nh//nkv,nparams)) | |
| o.write('/EMB ' + qrows(emb, V, dim) + ' def\n') | |
| for name, t, r, c in [('WQ',wq,dim,dim),('WK',wk,kvd,dim),('WV',wv,kvd,dim),('WO',wo,dim,dim),('W1',w1,hid,dim),('W2',w2,dim,hid),('W3',w3,hid,dim)]: | |
| o.write('/%s [\n' % name) | |
| for l in range(L): | |
| o.write(qrows(t[l*r*c:(l+1)*r*c], r, c) + '\n') | |
| o.write('] def\n') | |
| o.write('/RMSATT [' + ' '.join(realarr(rms_att[l*dim:(l+1)*dim]) for l in range(L)) + '] def\n') | |
| o.write('/RMSFFN [' + ' '.join(realarr(rms_ffn[l*dim:(l+1)*dim]) for l in range(L)) + '] def\n') | |
| o.write('/RMSF ' + realarr(rms_f) + ' def\n') | |
| # tokenizer: decode pieces into raw bytes | |
| tb = open(tokf,'rb').read(); q = 4; pieces=[] | |
| for i in range(V): | |
| sc, ln = struct.unpack('fi', tb[q:q+8]); q += 8 | |
| s = tb[q:q+ln]; q += ln | |
| if len(s)==6 and s.startswith(b'<0x') and s.endswith(b'>'): | |
| s = bytes([int(s[3:5],16)]) | |
| pieces.append(s) | |
| o.write('/VOCAB ' + psarr(['<'+s.hex()+'>' for s in pieces]) + ' def\n') | |
| o.close() | |
| print('params', nparams, 'dim', dim, 'hid', hid, 'L', L, 'nh', nh, 'nkv', nkv, 'V', V, 'seq', seq, 'used floats', p, 'of', len(f)) | |