Spaces:
Running
Running
| import { loadRuntime, encodePrompt, renderStory, MODEL_ID, REVISION } from './runtime.js'; | |
| import { sample, seededRandom, validateSettings } from './sampling.js'; | |
| let runtime, loading = false, running = false, stopped = false; | |
| const send = (type, data = {}) => postMessage({ type, ...data }); | |
| const yieldToEvents = () => new Promise(resolve => setTimeout(resolve, 0)); | |
| self.onmessage = async ({ data }) => { | |
| try { | |
| if (data.type === 'stop') { stopped = true; return; } | |
| if (data.type === 'load') { | |
| if (running || loading) return; | |
| loading = true; | |
| runtime?.decoder.past.dispose(); | |
| await runtime?.model.dispose(); | |
| runtime = null; | |
| try { | |
| runtime = await loadRuntime(data.device, info => send('loading', { info })); | |
| send('ready', { device: runtime.device, fallback: runtime.fallback, vocabulary: runtime.vocabulary, provenance: runtime.provenance }); | |
| } finally { loading = false; } | |
| return; | |
| } | |
| if (!runtime) throw new Error('Load the model first.'); | |
| if (data.type === 'inspect') { | |
| if (!running) send('inspection', { ...encodePrompt(runtime.tokenizer, data.prompt), prompt: data.prompt }); | |
| return; | |
| } | |
| if (data.type !== 'generate' || running) return; | |
| validateSettings(data.settings); | |
| const prompt = encodePrompt(runtime.tokenizer, data.prompt); | |
| if (prompt.ids.some(id => [0, 1, 2].includes(id))) throw new Error('Leave start, end and pad tokens out of your story opening.'); | |
| const input = [1, ...prompt.ids]; | |
| if (input.length >= 2048) throw new Error('Your opening fills the 2,048-token context. Shorten it to leave room for a continuation.'); | |
| running = true; | |
| stopped = false; | |
| const { decoder, tokenizer, device, provenance } = runtime; | |
| decoder.reset(); | |
| const generated = [], random = seededRandom(data.settings.seed); | |
| const started = performance.now(); | |
| let logits, firstTokenMs = null, reason = 'token_limit'; | |
| const emit = (type) => { | |
| const elapsedMs = performance.now() - started; | |
| const story = renderStory(tokenizer, [...prompt.ids, ...generated]); | |
| send(type, { | |
| ...story, ids: [...generated], promptIds: prompt.ids, unknown: prompt.unknown, | |
| settings: data.settings, prompt: data.prompt, device, reason, | |
| elapsedMs, firstTokenMs, tokens: generated.length, | |
| tokensPerSecond: generated.length ? generated.length / (elapsedMs / 1000) : 0, | |
| contextUsed: input.length + generated.length, | |
| model: { id: MODEL_ID, revision: REVISION, ...provenance }, | |
| }); | |
| }; | |
| try { | |
| for (let i = 0; i < input.length; i++) { | |
| if (stopped) break; | |
| logits = await decoder.forward(input[i]); | |
| if (i % 8 === 0) await yieldToEvents(); | |
| } | |
| const limit = Math.min(data.settings.maxTokens, 2048 - input.length); | |
| let lastUpdate = 0; | |
| for (let i = 0; i < limit && !stopped; i++) { | |
| const token = sample(logits, data.settings, random); | |
| generated.push(token); | |
| firstTokenMs ??= performance.now() - started; | |
| if (token === 2) { reason = 'eos'; break; } | |
| if (performance.now() - lastUpdate > 40) { | |
| emit('update'); | |
| lastUpdate = performance.now(); | |
| await yieldToEvents(); | |
| } | |
| // No extra forward call is needed after the last sampled token. | |
| if (i + 1 < limit && !stopped) logits = await decoder.forward(token); | |
| } | |
| if (stopped) reason = 'stopped'; | |
| else if (reason !== 'eos' && input.length + generated.length === 2048) reason = 'context_limit'; | |
| emit('done'); | |
| } finally { running = false; decoder.reset(); } | |
| } catch (error) { | |
| send('error', { message: error.message ?? String(error) }); | |
| } | |
| }; | |