Spaces:
Running
Running
| import { numpy as np, nn, jit, valueAndGrad, tree, blockUntilReady } from '@jax-js/jax'; | |
| import {sampleCharacter} from './sampling.js'; | |
| export function rng(seed = 42) { | |
| let a = seed >>> 0; | |
| return () => { a += 0x6D2B79F5; let t = a; t = Math.imul(t ^ t >>> 15, t | 1); t ^= t + Math.imul(t ^ t >>> 7, t | 61); return ((t ^ t >>> 14) >>> 0) / 4294967296; }; | |
| } | |
| export function tokenize(text) { | |
| const chars = [...new Set([...text])].sort(); | |
| if (chars.length < 2 || chars.length > 128) throw new Error('The dataset must contain 2 to 128 distinct characters.'); | |
| const index = new Map(chars.map((c, i) => [c, i])); | |
| const letters=[...text],paragraphStarts=[0]; | |
| for(let i=1;i<letters.length-1;i++)if(letters[i-1]==='\n'&&letters[i]==='\n'&&letters[i+1]!=='\n')paragraphStarts.push(i+1); | |
| return { chars, paragraphStarts, ids: [...text].map(c => index.get(c)), encode: s => [...s].map(c => { if (!index.has(c)) throw new Error(`Character missing from the dataset: ${JSON.stringify(c)}`); return index.get(c); }) }; | |
| } | |
| export function initParams(c, seed = 42) { | |
| const random = rng(seed); | |
| const weight = (shape, scale) => np.array(Float32Array.from({ length: shape.reduce((a,b) => a*b,1) }, () => (random()*2-1)*scale)).reshape(shape); | |
| if(c.architecture==='nanogpt') { | |
| // NanoGPT bias=False, dropout=0: LayerNorm, exact GELU, tied embedding/head. | |
| // Same distribution and residual scaling; JS/Python RNG streams differ. | |
| const normal=(shape,std=.02)=>np.array(Float32Array.from({length:shape.reduce((a,b)=>a*b,1)},()=>Math.sqrt(-2*Math.log(Math.max(1e-9,random())))*Math.cos(2*Math.PI*random())*std)).reshape(shape); | |
| const p={embed:normal([c.vocab,c.dim]),pos:normal([c.context,c.dim]),norm:np.ones([c.dim])}; | |
| for(let l=0;l<c.layers;l++) { | |
| for(const k of ['q','k','v'])p[`${l}.${k}`]=normal([c.dim,c.dim]); | |
| p[`${l}.o`]=normal([c.dim,c.dim],.02/Math.sqrt(2*c.layers)); | |
| p[`${l}.up`]=normal([c.dim,4*c.dim]); | |
| p[`${l}.down`]=normal([4*c.dim,c.dim],.02/Math.sqrt(2*c.layers)); | |
| p[`${l}.n1`]=np.ones([c.dim]);p[`${l}.n2`]=np.ones([c.dim]); | |
| } | |
| return p; | |
| } | |
| if(c.architecture==='modern') { | |
| // Modded-inspired FP32 variant: Gaussian fan-in init, zero residual projections, | |
| // RMSNorm, QK norm, RoPE, squared ReLU, soft-capped untied head. AdamW below. | |
| const normal=(shape,std)=>np.array(Float32Array.from({length:shape.reduce((a,b)=>a*b,1)},()=>Math.sqrt(-2*Math.log(Math.max(1e-9,random())))*Math.cos(2*Math.PI*random())*std)).reshape(shape); | |
| const zeroOrNormal=shape=>c.nonzeroInit?normal(shape,0.1):np.zeros(shape); | |
| const p={embed:normal([c.vocab,c.dim],1),norm:np.ones([c.dim]),head:zeroOrNormal([c.dim,c.vocab])}; | |
| for(let l=0;l<c.layers;l++) { | |
| for(const k of ['q','k','v'])p[`${l}.${k}`]=normal([c.dim,c.dim],Math.sqrt(.33/c.dim)); | |
| p[`${l}.o`]=zeroOrNormal([c.dim,c.dim]); | |
| p[`${l}.up`]=normal([c.dim,4*c.dim],Math.sqrt(.33/c.dim)); | |
| p[`${l}.down`]=zeroOrNormal([4*c.dim,c.dim]); | |
| p[`${l}.n1`]=np.ones([c.dim]);p[`${l}.n2`]=np.ones([c.dim]); | |
| } | |
| return p; | |
| } | |
| const p = { embed: weight([c.vocab,c.dim],0.2), pos: weight([c.context,c.dim],0.02) }; | |
| for(let l=0;l<c.layers;l++) { | |
| for(const k of ['q','k','v','o']) p[`${l}.${k}`] = weight([c.dim,c.dim],Math.sqrt(3/c.dim)); | |
| p[`${l}.up`] = weight([c.dim,4*c.dim],Math.sqrt(3/c.dim)); | |
| p[`${l}.down`] = weight([4*c.dim,c.dim],Math.sqrt(3/(4*c.dim))); | |
| p[`${l}.n1`] = np.ones([c.dim]); p[`${l}.n2`] = np.ones([c.dim]); | |
| } | |
| p.norm = np.ones([c.dim]); p.head = weight([c.dim,c.vocab],0.02); | |
| // Keep the same RNG draws for other weights, for the positional ablation. | |
| if(c.architecture==='rotary'){p.pos.dispose();delete p.pos;} | |
| return p; | |
| } | |
| const norm = (x,w) => x.ref.div(np.sqrt(np.mean(np.square(x),-1,{keepdims:true}).add(1e-5))).mul(w); | |
| const layerNorm=(x,w)=>{const centered=x.ref.sub(np.mean(x,-1,{keepdims:true}));return norm(centered,w);}; | |
| const rms = x => x.ref.div(np.sqrt(np.mean(np.square(x),-1,{keepdims:true}).add(1e-5))); | |
| function rotary(x,t,hd,full=false) { | |
| const angles=Array.from({length:t*hd/2},(_,i)=>{ | |
| const j=i%(hd/2),pos=Math.floor(i/(hd/2)); | |
| return full?pos*(1/10000)**(2*j/hd):j<hd/4?pos*(1/1024)**(j/Math.max(1,hd/4-1)):0; | |
| }); | |
| const cs=np.array(angles.map(Math.cos)).reshape([1,t,1,hd/2]); | |
| const sn=np.array(angles.map(Math.sin)).reshape([1,t,1,hd/2]); | |
| const [a,b]=np.split(x,2,-1); | |
| return np.concatenate([a.ref.mul(cs.ref).add(b.ref.mul(sn.ref)),b.mul(cs).sub(a.mul(sn))],-1); | |
| } | |
| // Each argument is consumed exactly once, including every parameter leaf. | |
| export function forward(p, ids, c) { | |
| const [b,t] = ids.shape, d = c.dim, h = c.heads; | |
| const modern=c.architecture==='modern',rope=c.architecture==='rotary',nano=c.architecture==='nanogpt',normalize=nano?layerNorm:norm; | |
| let x = np.matmul(nn.oneHot(ids,c.vocab),nano?p.embed.ref:p.embed); | |
| if(!modern&&!rope)x=x.add(t===c.context?p.pos:p.pos.slice([0,t],[])); | |
| for(let l=0;l<c.layers;l++) { | |
| const n = normalize(x.ref,p[`${l}.n1`]); | |
| let q = np.matmul(n.ref,p[`${l}.q`]).reshape([b,t,h,d/h]); | |
| let k = np.matmul(n.ref,p[`${l}.k`]).reshape([b,t,h,d/h]); | |
| const v = np.matmul(n,p[`${l}.v`]).reshape([b,t,h,d/h]); | |
| if(modern){q=rotary(rms(q),t,d/h);k=rotary(rms(k),t,d/h);} | |
| if(rope){q=rotary(q,t,d/h,true);k=rotary(k,t,d/h,true);} | |
| const a = nn.dotProductAttention(q,k,v,{isCausal:true,...(modern?{scale:.12}:{})}).reshape([b,t,d]); | |
| x = x.add(np.matmul(a,p[`${l}.o`])); | |
| const z = normalize(x.ref,p[`${l}.n2`]); | |
| const u=np.matmul(z,p[`${l}.up`]); | |
| x = x.add(np.matmul(modern?np.square(nn.relu(u)):nn.gelu(u,{approximate:!nano}),p[`${l}.down`])); | |
| } | |
| const logits=np.matmul(normalize(x,p.norm),nano?p.embed.transpose():p.head); | |
| return modern?logits.ref.mul(15).div(np.sqrt(np.square(logits).add(225))):logits; | |
| } | |
| export function loss(p,x,y,c) { | |
| return nn.logSoftmax(forward(p,x,c),-1).mul(nn.oneHot(y,c.vocab)).sum(-1).mean().neg(); | |
| } | |
| export function batch(ids,c,random) { | |
| const x = [], y = []; | |
| for(let b=0;b<c.batch;b++) { | |
| const start = Math.floor(random()*(ids.length-c.context)); | |
| x.push(...ids.slice(start,start+c.context)); y.push(...ids.slice(start+1,start+c.context+1)); | |
| } | |
| return [np.array(x,{dtype:np.int32}).reshape([c.batch,c.context]),np.array(y,{dtype:np.int32}).reshape([c.batch,c.context])]; | |
| } | |
| // Paragraph starts use Unicode code-point offsets, like token ids. | |
| // Circular continuation includes the last paragraph without padding fake targets. | |
| export function paragraphBatch(ids,c,random,starts) { | |
| if(!starts?.length)throw new Error('No paragraph starts available.'); | |
| const x=[],y=[]; | |
| for(let b=0;b<c.batch;b++){ | |
| const start=starts[Math.floor(random()*starts.length)]; | |
| for(let i=0;i<c.context;i++){x.push(ids[(start+i)%ids.length]);y.push(ids[(start+i+1)%ids.length]);} | |
| } | |
| return [np.array(x,{dtype:np.int32}).reshape([c.batch,c.context]),np.array(y,{dtype:np.int32}).reshape([c.batch,c.context])]; | |
| } | |
| export class Trainer { | |
| constructor(c,seed=42) { | |
| this.c=c; this.p=initParams(c,seed); this.step=0; | |
| this.m=tree.map(p=>np.zeros(p.shape),this.p); this.v=tree.map(p=>np.zeros(p.shape),this.p); | |
| this.lossGrad=jit(valueAndGrad((p,x,y)=>loss(p,x,y,c))); | |
| this.infer=jit((p,x)=>forward(p,x,c)); | |
| this.inferLast=jit((p,x,i)=>np.take(forward(p,x,c),i,1)); | |
| this.evaluate=jit((p,x,y)=>nn.logSoftmax(forward(p,x,c),-1).mul(nn.oneHot(y,c.vocab)).sum(-1).slice([], [Math.max(0,c.context-16),c.context]).mean().neg()); | |
| this.beta2=c.beta2??.999; | |
| this.update=jit((p,m,v,g,lr,b1,b2)=>{ | |
| // Global norm clipping before Adam. | |
| const total=Object.values(g).map(x=>np.square(x.ref).sum()).reduce((a,b)=>a.add(b)); | |
| const scale=np.minimum(1,np.reciprocal(np.sqrt(total).add(1e-8))); | |
| const nextP={},nextM={},nextV={}; | |
| for(const k of Object.keys(p)) { | |
| const grad=g[k].mul(scale.ref); | |
| const mm=m[k].mul(0.9).add(grad.ref.mul(0.1)); | |
| const vv=v[k].mul(this.beta2).add(np.square(grad).mul(1-this.beta2)); | |
| const decayed=p[k].mul(np.array(1).sub(lr.ref.mul(p[k].shape.length===2?(c.weightDecay??0):0))); | |
| nextP[k]=decayed.sub(mm.ref.div(b1.ref).div(np.sqrt(vv.ref.div(b2.ref)).add(1e-8)).mul(lr.ref)); | |
| nextM[k]=mm; nextV[k]=vv; | |
| } | |
| scale.dispose(); lr.dispose(); b1.dispose(); b2.dispose(); | |
| return [nextP,nextM,nextV]; | |
| }); | |
| } | |
| async train(x,y,rate=0.003) { | |
| const [l,g]=this.lossGrad(tree.ref(this.p),x,y); | |
| this.step++; | |
| [this.p,this.m,this.v]=this.update(this.p,this.m,this.v,g,np.array(rate),np.array(1-0.9**this.step),np.array(1-this.beta2**this.step)); | |
| await blockUntilReady([this.p,this.m,this.v]); | |
| const value=(await l.data())[0]; | |
| if(!Number.isFinite(value)) throw new Error('Non-finite loss. Reset and lower the learning rate.'); | |
| return value; | |
| } | |
| async generate(prefix,tokenizer,{length=80,temperature=0.8,minP=0,seed=123,bucketed=this.c.fastSampling??false,cancel=()=>false,onToken=()=>{}}={}) { | |
| const ids=tokenizer.encode(prefix); if(!ids.length) throw new Error('The prompt must not be empty.'); | |
| const random=rng(seed); let text=''; | |
| for(let i=0;i<length&&!cancel();i++) { | |
| const ctx=ids.slice(-this.c.context), pos=ctx.length-1; | |
| const size=bucketed?Math.min(this.c.context,Math.max(16,2**Math.ceil(Math.log2(ctx.length)))):this.c.context; | |
| const padded=[...ctx,...Array(size-ctx.length).fill(0)]; | |
| const x=np.array(padded,{dtype:np.int32}).reshape([1,size]); | |
| const logits=bucketed?this.inferLast(tree.ref(this.p),x,np.array(pos,{dtype:np.int32})):this.infer(tree.ref(this.p),x); | |
| const raw=await logits.data(),values=bucketed?raw:raw.slice(pos*this.c.vocab,(pos+1)*this.c.vocab); | |
| const choice=sampleCharacter(values,temperature,random,minP); | |
| ids.push(choice); text+=tokenizer.chars[choice]; onToken(text); | |
| if(i%8===0) await new Promise(r=>setTimeout(r,0)); | |
| } | |
| return text; | |
| } | |
| dispose(){tree.dispose([this.p,this.m,this.v]); this.lossGrad.dispose(); this.infer.dispose(); this.inferLast.dispose(); this.evaluate.dispose(); this.update.dispose();} | |
| } | |