nx_nofloat_generate_gate.nx source

1// nx_nofloat_generate_gate.nx -- AUTOREGRESSIVE TEXT GENERATION from a trained no-float LM, integer Q16 2// (CAP-NF-GENERATE). Train an attention-only char-LM to memorize a short string, then argmax-decode it back: 3// seed the first char, run the model on the generated-so-far prefix, take the last position's argmax, append, 4// repeat. Reproducing the text = the model generates. Lean (nx_cc .s ceiling): 2 clm_fwd sites (train, gen). 5// expect_exit: 0 Sovereign: nx_nofloat_autograd + nx_syscalls. 6import "nx_nofloat_autograd.nx" 7import "nx_g_puts_lib.nx" 8import "nx_syscalls.nx" 9func g_pn(v: i64) -> i64 { let b: *u8=sys_mmap(28); var x: i64=v; if x<0{b[0]=45;sys_write(1,b,1);x=0-x} if x==0{b[0]=48;sys_write(1,b,1);return 0} var d: i64=0; var y: i64=x; while y>0{d=d+1;y=y/10} var i: i64=d-1; y=x; while i>=0{b[i]=(48+(y%10)) as u8;y=y/10;i=i-1} sys_write(1,b,d); return 0 } 10func slen(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} return n } 11func allmatch(gen: *i64, seq: *i64, L: i64) -> i64 { var ok: i64=1; var i: i64=0; while i<L { if gen[i]!=seq[i] { ok=0 } i=i+1 } return ok } 12func dini(a: *i64, n: i64, sd: i64) -> i64 { var i: i64=0; while i<n { a[i]=(((i*7+sd*13+1)%11)-5)*13107; i=i+1 } return 0 } 13func amx(tape: *i64, vals: *i64, logn: i64, r: i64, V: i64) -> i64 { let o: i64=tape[7*logn+5]; var b: i64=0; var bv: i64=vals[o+r*V]; var j: i64=1; while j<V { if vals[o+r*V+j]>bv { bv=vals[o+r*V+j]; b=j } j=j+1 } return b } 14func clm_fwd(tape: *i64, vals: *i64, st: *i64, W: *i64, ids: *i64, tgt: *i64, T: i64, dm: i64, V: i64, scale: i64, leaves: *i64) -> i64 { 15 let E: *i64=W[0] as *i64; let Wq: *i64=W[1] as *i64; let Wk: *i64=W[2] as *i64; let Wv: *i64=W[3] as *i64; let Wo: *i64=W[4] as *i64; let Wlm: *i64=W[5] as *i64 16 st[0]=0; st[1]=0 17 let nE: i64=nfa_leaf(tape,vals,st,V,dm,E,0) 18 let nWq: i64=nfa_leaf(tape,vals,st,dm,dm,Wq,0) 19 let nWk: i64=nfa_leaf(tape,vals,st,dm,dm,Wk,0) 20 let nWv: i64=nfa_leaf(tape,vals,st,dm,dm,Wv,0) 21 let nWo: i64=nfa_leaf(tape,vals,st,dm,dm,Wo,0) 22 let nWlm: i64=nfa_leaf(tape,vals,st,dm,V,Wlm,0) 23 let nX: i64=nfa_embed(tape,vals,st,nE,ids,T) 24 let nXn: i64=nfa_rmsnorm_rows(tape,vals,st,nX) 25 let nQ: i64=nfa_matmul(tape,vals,st,nXn,nWq) 26 let nK: i64=nfa_matmul(tape,vals,st,nXn,nWk) 27 let nV: i64=nfa_matmul(tape,vals,st,nXn,nWv) 28 let nQr: i64=nfa_rope(tape,vals,st,nQ) 29 let nKr: i64=nfa_rope(tape,vals,st,nK) 30 let nS: i64=nfa_matmul_nt(tape,vals,st,nQr,nKr) 31 let nSs: i64=nfa_cmul(tape,vals,st,nS,scale) 32 let nA: i64=nfa_softmax_rows(tape,vals,st,nSs,1) 33 let nO: i64=nfa_matmul(tape,vals,st,nA,nV) 34 let nOp: i64=nfa_matmul(tape,vals,st,nO,nWo) 35 let nH: i64=nfa_vadd(tape,vals,st,nX,nOp) 36 let nHn: i64=nfa_rmsnorm_rows(tape,vals,st,nH) 37 let nLg: i64=nfa_matmul(tape,vals,st,nHn,nWlm) 38 let nLoss: i64=nfa_softce_rows(tape,vals,st,nLg,tgt) 39 leaves[0]=nE; leaves[1]=nWq; leaves[2]=nWk; leaves[3]=nWv; leaves[4]=nWo; leaves[5]=nWlm; leaves[6]=nLg 40 return nLoss 41} 42func step_all(tape: *i64, grads: *i64, W: *i64, WN: *i64, leaves: *i64, nW: i64, lr: i64, clip: i64, gb: *i64) -> i64 { 43 var i: i64=0 44 while i<nW { let ar: *i64=W[i] as *i64; let cn: i64=WN[i]; let nd: i64=leaves[i]; var c: i64=0; while c<cn { var g: i64=nfa_grad(tape,grads,nd,c); if g>clip{g=clip} if g<0-clip{g=0-clip} gb[c]=g; c=c+1 } nfa_sgd(ar,gb,cn,lr); i=i+1 } 45 return 0 46} 47func main() -> i64 { 48 g_puts("nx_nofloat_generate (autoregressive decode from a trained no-float LM, Q16)\n" as *u8) 49 let corpus: *u8 = "abcdefgh" as *u8 50 let L: i64 = slen(corpus) 51 let c2i: *i64 = sys_mmap(256*8) as *i64 52 var iz: i64=0; while iz<256 { c2i[iz]=0-1; iz=iz+1 } 53 let i2c: *i64 = sys_mmap(256*8) as *i64 54 var V: i64=0; var pz: i64=0 55 while pz<L { let ch: i64=corpus[pz] as i64; if c2i[ch]<0 { c2i[ch]=V; i2c[V]=ch; V=V+1 } pz=pz+1 } 56 let seq: *i64 = sys_mmap(L*8) as *i64 57 pz=0; while pz<L { seq[pz]=c2i[corpus[pz] as i64]; pz=pz+1 } 58 let T: i64 = L-1 59 let ids: *i64 = sys_mmap(T*8) as *i64; let tgt: *i64 = sys_mmap(T*8) as *i64 60 pz=0; while pz<T { ids[pz]=seq[pz]; tgt[pz]=seq[pz+1]; pz=pz+1 } 61 let dm: i64=16; let scale: i64=16384 62 let tape: *i64 = sys_mmap(512*7*8) as *i64 63 let vals: *i64 = sys_mmap(65536*8) as *i64 64 let grads: *i64 = sys_mmap(65536*8) as *i64 65 let st: *i64 = sys_mmap(2*8) as *i64 66 let nW: i64=6 67 let W: *i64 = sys_mmap(nW*8) as *i64 68 let WN: *i64 = sys_mmap(nW*8) as *i64 69 WN[0]=V*dm; WN[1]=dm*dm; WN[2]=dm*dm; WN[3]=dm*dm; WN[4]=dm*dm; WN[5]=dm*V 70 var wi: i64=0 71 while wi<nW { let a: *i64=sys_mmap(WN[wi]*8) as *i64; dini(a,WN[wi],wi+1); W[wi]=a as i64; wi=wi+1 } 72 let leaves: *i64 = sys_mmap(8*8) as *i64 73 let gbuf: *i64 = sys_mmap(4096*8) as *i64 74 var ep: i64=0 75 while ep < 15000 { let nl: i64=clm_fwd(tape,vals,st,W,ids,tgt,T,dm,V,scale,leaves); nfa_backward(tape,vals,grads,st[0],nl); step_all(tape,grads,W,WN,leaves,nW,6554,262144,gbuf); ep=ep+1 } 76 let gen: *i64 = sys_mmap(L*8) as *i64 77 var zz: i64=0; while zz<T { gen[zz]=0; zz=zz+1 } 78 gen[0]=seq[0] 79 var sp: i64=1 80 while sp <= T { let ng: i64=clm_fwd(tape,vals,st,W,gen,ids,T,dm,V,scale,leaves); gen[sp]=amx(tape,vals,leaves[6],sp-1,V); sp=sp+1 } 81 g_puts(" target : " as *u8); g_puts(corpus); g_puts("\n GENERATED: " as *u8) 82 var gp: i64=0; while gp<L { let cb: *u8=sys_mmap(2); cb[0]=i2c[gen[gp]] as u8; sys_write(1,cb,1); gp=gp+1 } 83 g_puts("\n" as *u8) 84 let am: i64 = allmatch(gen, seq, L) 85 g_puts(" reproduced=" as *u8); g_pn(am); g_puts(" (1=GREEN: generation reproduces the trained text)\n" as *u8) 86 sys_exit(1 - am) 87 return 0 88}