Spaces:
Running
Running
| import { init, defaultDevice, getWebGPUDevice } from '@jax-js/jax'; | |
| import { Trainer, tokenize, batch, paragraphBatch, rng } from './model.js'; | |
| import {validateExploration} from '../public/exploration.js'; | |
| let exploring=false; | |
| let model, tok, config, random, running=false, busy=false, stopped=false, elapsed=0, request,initialSampleComplete=false; | |
| const post=(type,data={})=>self.postMessage({type,...data}); | |
| const yieldTask=()=>new Promise(r=>setTimeout(r,0)); | |
| function fail(e){ running=false; stopped=true; post('error',{message:e?.message||String(e)}); } | |
| async function sample(){ | |
| const step=model.step;post('sampling',{step}); | |
| const text=await model.generate(request.prefix,tok,{length:request.sampleLength??80,temperature:request.temperature,seed:123,cancel:()=>!running||stopped,onToken:text=>{if([...text].length%8===0)post('sample-progress',{step,prefix:request.prefix,text});}}); | |
| if(step===0)initialSampleComplete=[...text].length===(request.sampleLength??80); | |
| post('sample',{step,prefix:request.prefix,text}); | |
| } | |
| async function loop(){ | |
| if(busy||stopped)return;busy=true; | |
| try { | |
| if(running&&model.step===0&&!initialSampleComplete)await sample(); | |
| while(running&&!stopped){ | |
| const start=performance.now(); | |
| const value=await model.train(...(request.sampling==='paragraphs'?paragraphBatch(tok.ids,config,random,tok.paragraphStarts):batch(tok.ids,config,random)),request.rate*(request.schedule==='cooldown'?Math.min(1,(1-model.step/request.maxSteps)/.7):1)); | |
| const ms=performance.now()-start;elapsed+=ms; | |
| post('step',{step:model.step,loss:value,ms,elapsed,tokens:model.step*config.batch*config.context}); | |
| if(model.step%4===0)await yieldTask(); | |
| if(running&&(model.step%request.sampleEvery===0||model.step>=request.maxSteps)) await sample(); | |
| if(model.step>=request.maxSteps){running=false;post('complete');} | |
| } | |
| }catch(e){fail(e);}finally{busy=false;if(!stopped)post('paused',{step:model.step});} | |
| } | |
| self.onmessage=async({data})=>{ | |
| if(data.type==='cancel-explore'){exploring=false;return;} | |
| if(data.type==='explore'){ | |
| if(!model||busy||running||stopped||model.step<request.maxSteps){post('explore-error',{id:data.id,message:'Wait for training or the current generation to finish.'});return;} | |
| busy=true;const started=performance.now(); | |
| try{ | |
| const options=validateExploration(data,request.text);exploring=true; | |
| const text=await model.generate(options.prefix,tok,{...options,cancel:()=>!exploring||stopped,onToken:text=>{if([...text].length%8===0)post('explore-progress',{id:data.id,text});}}); | |
| post('explore-result',{id:data.id,...options,text,step:model.step,ms:performance.now()-started,cancelled:!exploring}); | |
| }catch(e){post('explore-error',{id:data.id,message:e.message||String(e)});} | |
| finally{exploring=false;busy=false;} | |
| return; | |
| } | |
| if(data.type==='pause'){running=false;return;} | |
| if(data.type==='resume'){if(!model||stopped||busy||model.step>=request.maxSteps)return;running=true;post('running');void loop();return;} | |
| if(data.type!=='start'||busy||model)return; | |
| busy=true; | |
| try { | |
| const loadStart=performance.now();request=data;tok=tokenize(data.text);tok.encode(data.prefix); | |
| if(!data.prefix.length)throw new Error('Enter a prompt.'); | |
| if(tok.ids.length<129||tok.ids.length>100000)throw new Error('The dataset must contain between 129 and 100,000 characters.'); | |
| post('status',{message:'Initializing the backend…'}); | |
| let available=[];let fallback=''; | |
| if(data.backend!=='wasm') { | |
| try{available=await init('webgpu');}catch(e){fallback=e.message;} | |
| } | |
| const backend=available.includes('webgpu')?'webgpu':'wasm'; | |
| if(data.backend==='webgpu'&&backend!=='webgpu')throw new Error(self.isSecureContext?'WebGPU unavailable. Select Auto or WASM.':'WebGPU unavailable over HTTP: open the demo over HTTPS or select Auto / WASM.'); | |
| if(backend==='wasm')await init('wasm'); | |
| defaultDevice(backend); | |
| if(backend==='webgpu')getWebGPUDevice().lost.then(info=>{if(!stopped)fail(new Error(`GPU device lost: ${info.message}. Reset in WASM mode.`));}); | |
| const large=data.size==='standard'; | |
| config={architecture:['nanogpt','rotary'].includes(data.architecture)?data.architecture:undefined,fastSampling:data.fastSampling===true&&backend==='wasm',vocab:tok.chars.length,dim:large?64:32,layers:large?2:1,heads:4,context:[32,64,128,256].includes(data.context)?data.context:32,batch:[4,8].includes(data.batch)?data.batch:4,beta2:data.optimizer==='tuned'?.95:.999,weightDecay:data.optimizer==='tuned'?.01:0}; | |
| if(tok.ids.length<=config.context)throw new Error('The dataset must be longer than the selected context.'); | |
| if(data.sampling==='paragraphs'&&tok.paragraphStarts.length<2)throw new Error('This recipe requires at least two paragraphs separated by a blank line.'); | |
| random=rng(99);model=new Trainer(config,42); | |
| const params=Object.values(model.p).reduce((s,p)=>s+p.size,0); | |
| post('ready',{loadMs:performance.now()-loadStart,backend,params,vocab:tok.chars.length,config,fallback:backend==='wasm'&&data.backend==='auto'?(self.isSecureContext?'WebGPU unavailable; using WASM on CPU. ':'HTTP connection: using WASM on CPU. WebGPU requires HTTPS. ')+fallback:''}); | |
| running=true;post('status',{message:'First generation and compilation…'});await sample(); | |
| busy=false;if(running){post('running');void loop();}else post('paused',{step:0}); | |
| }catch(e){busy=false;fail(e);} | |
| }; | |