Spaces:
Running on Zero
Running on Zero
| import { Wllama } from '@wllama/wllama'; | |
| import wasmURL from '@wllama/wllama/esm/wasm/wllama.wasm?url'; | |
| import { createSHA256 } from 'hash-wasm'; | |
| import { assertModel, recentMessages } from './model-utils.js'; | |
| const folderName = 'cabin-local-models-v1'; | |
| const quiet = {log(){},debug(){},warn(...v){console.warn(...v)},error(...v){console.error(...v)}}; | |
| const abortCheck = signal => {if (signal?.aborted) throw new DOMException('Operation cancelled', 'AbortError');}; | |
| async function directory() { | |
| if (!navigator.storage?.getDirectory) throw new Error('Model storage is unavailable. Uncheck “Keep the weights” to try in memory, or use a recent browser outside private browsing.'); | |
| return (await navigator.storage.getDirectory()).getDirectoryHandle(folderName,{create:true}); | |
| } | |
| export async function cachedModels(models) { | |
| const ids=[]; | |
| try {const dir=await directory();for(const model of models){try{const f=await (await dir.getFileHandle(model.id)).getFile();if(f.size===model.bytes)ids.push(model.id);}catch{}}}catch{} | |
| return ids; | |
| } | |
| export async function deleteModel(model) { | |
| const dir=await directory();try{await dir.removeEntry(model.id);}catch(e){if(e.name!=='NotFoundError')throw e;} | |
| } | |
| export class LocalEngine { | |
| constructor() { this.runtime=null; this.model=null; this.history=[]; } | |
| async unload() { | |
| const old=this.runtime;this.runtime=null;this.model=null;this.history=[]; | |
| if(old)await old.exit(); | |
| } | |
| async load(model, signal, onProgress, persist=true) { | |
| await this.unload(); abortCheck(signal); | |
| const dir=persist?await directory():null; let file; | |
| if(dir)try {file=await(await dir.getFileHandle(model.id)).getFile();}catch{} | |
| const hash=await createSHA256();hash.init(); | |
| if(file?.size===model.bytes){ | |
| onProgress({phase:'verify',loaded:0,total:model.bytes}); | |
| const reader=file.stream().getReader();let loaded=0; | |
| try {while(true){abortCheck(signal);const {done,value}=await reader.read();if(done)break;hash.update(value);loaded+=value.byteLength;onProgress({phase:'verify',loaded,total:model.bytes});}} | |
| finally{await reader.cancel();} | |
| try {assertModel(model,file.size,hash.digest());}catch(e){await dir.removeEntry(model.id);throw e;} | |
| } else { | |
| const estimate=await navigator.storage?.estimate?.().catch(()=>({}))||{}; | |
| if(persist&&estimate.quota && estimate.quota-estimate.usage < model.bytes*1.1)throw new Error('Not enough storage. Uncheck “Keep the weights” to try in memory, or free some storage.'); | |
| const response=await fetch(model.downloadURL,{signal,credentials:'omit'}); | |
| if(!response.ok || !response.body)throw new Error(`Download unavailable (HTTP ${response.status}). Try again.`); | |
| let handle,writer;const chunks=[]; | |
| if(dir){try{handle=await dir.getFileHandle(model.id,{create:true});writer=await handle.createWritable();}catch(e){await response.body.cancel().catch(()=>{});await dir.removeEntry(model.id).catch(()=>{});throw e;}} | |
| const reader=response.body.getReader();let loaded=0; | |
| try { | |
| while(true){abortCheck(signal);const {done,value}=await reader.read();if(done)break;loaded+=value.byteLength;if(loaded>model.bytes)throw new Error('Unexpected model file size.');hash.update(value);if(writer)await writer.write(value);else chunks.push(value);onProgress({phase:'download',loaded,total:model.bytes});} | |
| assertModel(model,loaded,hash.digest());if(writer){await writer.close();file=await handle.getFile();}else{file=new Blob(chunks);chunks.length=0;} | |
| }catch(e){if(writer)await writer.abort().catch(()=>{});if(dir)await dir.removeEntry(model.id).catch(()=>{});throw e;} | |
| finally{await reader.cancel().catch(()=>{});} | |
| } | |
| abortCheck(signal);onProgress({phase:'load',loaded:1,total:1}); | |
| this.runtime=new Wllama({default:wasmURL},{suppressNativeLog:true,logger:quiet}); | |
| try { | |
| await this.runtime.loadModel([file],{n_ctx:2048,n_batch:128,n_ubatch:128,n_threads:Math.min(4,navigator.hardwareConcurrency||2),n_gpu_layers:0,offload_kqv:false,no_kv_offload:true,seed:42}); | |
| abortCheck(signal);this.model=model; | |
| }catch(e){await this.unload().catch(()=>{});throw e;} | |
| } | |
| async generate(question, signal, onText) { | |
| if(!this.runtime||!this.model)throw new Error('Load a model before starting.'); | |
| let answer=''; | |
| await this.runtime.createChatCompletion({ | |
| messages:recentMessages(this.history,question),stream:true,abortSignal:signal, | |
| max_tokens:160,temperature:0.6,top_k:40,top_p:0.9,seed:42,cache_prompt:false, | |
| onData:chunk=>{const text=chunk.choices?.[0]?.delta?.content;if(text){answer+=text;onText(answer);}} | |
| }); | |
| abortCheck(signal); | |
| if(!answer.trim())throw new Error('The model did not produce a response. Try a shorter question.'); | |
| this.history=recentMessages(this.history,question).slice(1).concat({role:'assistant',content:answer}); | |
| return answer; | |
| } | |
| info(){return {adapter:'wllama 3.13.0',llamaVersion:Wllama.getLibllamaVersion(),model:this.model?.id||null,backend:'WebAssembly CPU',threads:this.runtime?.getNumThreads()||null,crossOriginIsolated:globalThis.crossOriginIsolated===true};} | |
| } | |