Spaces:
Running
Running
| // Inference worker: downloads the brain once, then runs one recurrent step per request. | |
| // | |
| // The main thread drives the pace so the text appears as the fly says it, rather than | |
| // racing ahead. Generation is cheap enough (tens of tokens per second) that the | |
| // bottleneck is speech, not the model. | |
| import { FlyEngine } from './engine.js'; | |
| let engine = null; | |
| let manifest = null; | |
| let cloudIndex = null; | |
| let rasterIndex = null; | |
| let logits = null; | |
| let generated = 0; | |
| const TENSORS = ['wte', 'in_proj', 'gain', 'rec_gain', 'bias', 'ln_weight', 'ln_bias', 'head_a', 'head_b']; | |
| async function fetchTracked(url, label, onBytes) { | |
| const response = await fetch(url); | |
| if (!response.ok) throw new Error(`${label}: HTTP ${response.status}`); | |
| const total = Number(response.headers.get('content-length')) || 0; | |
| if (!response.body) return { buffer: await response.arrayBuffer(), total }; | |
| const reader = response.body.getReader(); | |
| const chunks = []; | |
| let received = 0; | |
| for (;;) { | |
| const { done, value } = await reader.read(); | |
| if (done) break; | |
| chunks.push(value); | |
| received += value.length; | |
| onBytes(value.length, label); | |
| } | |
| const merged = new Uint8Array(received); | |
| let at = 0; | |
| for (const chunk of chunks) { | |
| merged.set(chunk, at); | |
| at += chunk.length; | |
| } | |
| return { buffer: merged.buffer, total: received }; | |
| } | |
| function pickIndices(valid, count, neurons) { | |
| // Even stride over neurons that have a measured soma position, so the sample is | |
| // spread across the brain instead of clustered at low indices. | |
| const usable = []; | |
| for (let i = 0; i < neurons; i += 1) if (valid[i]) usable.push(i); | |
| const step = usable.length / count; | |
| const out = new Int32Array(count); | |
| for (let i = 0; i < count; i += 1) out[i] = usable[Math.min(usable.length - 1, Math.floor(i * step))]; | |
| return out; | |
| } | |
| async function load({ modelBase, graphBase, cloudCount, rasterCount }) { | |
| let done = 0; | |
| const report = (bytes, label) => { | |
| done += bytes; | |
| self.postMessage({ type: 'progress', bytes: done, label }); | |
| }; | |
| manifest = await (await fetch(`${modelBase}/manifest.json`)).json(); | |
| const config = manifest.config; | |
| const params = {}; | |
| for (const name of TENSORS) { | |
| const { buffer } = await fetchTracked(`${modelBase}/${name}.f32`, name, report); | |
| params[name] = new Float32Array(buffer); | |
| } | |
| const [offsets, source, weight, inIndex, soma, valid] = await Promise.all([ | |
| fetchTracked(`${graphBase}/graph/edges_offsets.i32`, 'edges', report), | |
| fetchTracked(`${graphBase}/graph/edges_source.u16`, 'edges', report), | |
| fetchTracked(`${graphBase}/graph/edges_weight.f32`, 'synapses', report), | |
| fetchTracked(`${graphBase}/interface/in_index.i32`, 'interface', report), | |
| fetchTracked(`${graphBase}/neurons/soma_position.i32`, 'anatomy', report), | |
| fetchTracked(`${graphBase}/neurons/soma_valid.u8`, 'anatomy', report), | |
| ]); | |
| const graph = { | |
| offsets: new Int32Array(offsets.buffer), | |
| source: new Uint16Array(source.buffer), | |
| weight: new Float32Array(weight.buffer), | |
| inIndex: new Int32Array(inIndex.buffer), | |
| }; | |
| if (graph.offsets.length !== config.neurons + 1) throw new Error('CSR offsets do not match the model'); | |
| if (graph.source.length !== graph.weight.length) throw new Error('edge arrays disagree'); | |
| const positions = new Int32Array(soma.buffer); | |
| const somaValid = new Uint8Array(valid.buffer); | |
| cloudIndex = pickIndices(somaValid, cloudCount, config.neurons); | |
| rasterIndex = cloudIndex.filter((_value, i) => i % Math.floor(cloudCount / rasterCount) === 0) | |
| .slice(0, rasterCount); | |
| const cloudPositions = new Float32Array(cloudIndex.length * 3); | |
| for (let i = 0; i < cloudIndex.length; i += 1) { | |
| const at = cloudIndex[i] * 3; | |
| cloudPositions[i * 3] = positions[at]; | |
| cloudPositions[i * 3 + 1] = positions[at + 1]; | |
| cloudPositions[i * 3 + 2] = positions[at + 2]; | |
| } | |
| engine = new FlyEngine({ params, graph, config }); | |
| self.postMessage({ | |
| type: 'ready', | |
| manifest: { | |
| experiment: manifest.experiment, | |
| selector: manifest.selector, | |
| updates: manifest.updates, | |
| parameters: manifest.trainable_parameters, | |
| edges: graph.weight.length, | |
| neurons: config.neurons, | |
| eos: config.eos_token_id, | |
| bos: config.bos_token_id, | |
| }, | |
| cloudPositions, | |
| rasterCount: rasterIndex.length, | |
| }, [cloudPositions.buffer]); | |
| } | |
| function sample() { | |
| const state = engine.state; | |
| const cloud = new Float32Array(cloudIndex.length); | |
| for (let i = 0; i < cloudIndex.length; i += 1) cloud[i] = Math.abs(state[cloudIndex[i]]); | |
| const raster = new Float32Array(rasterIndex.length); | |
| for (let i = 0; i < rasterIndex.length; i += 1) raster[i] = Math.abs(state[rasterIndex[i]]); | |
| // tanh saturates, so a peak magnitude is pinned near 1 and says nothing. Count how | |
| // many neurons are strongly driven instead; that number actually moves. | |
| let active = 0; | |
| for (let i = 0; i < state.length; i += 1) { | |
| const value = state[i] < 0 ? -state[i] : state[i]; | |
| if (value > 0.5) active += 1; | |
| } | |
| return { cloud, raster, active }; | |
| } | |
| function prime(promptIds) { | |
| engine.reset(); | |
| generated = 0; | |
| for (const id of promptIds) logits = engine.step(id); | |
| const { cloud, raster, active } = sample(); | |
| self.postMessage({ type: 'primed', meanAbs: engine.meanAbsState, active, cloud, raster }, | |
| [cloud.buffer, raster.buffer]); | |
| } | |
| function next({ temperature, topP }) { | |
| const started = performance.now(); | |
| const probabilities = engine.probabilities(logits); | |
| const token = engine.pick(logits, { temperature, topP }); | |
| const probability = probabilities[token]; | |
| let best = 0; | |
| for (let v = 1; v < probabilities.length; v += 1) if (probabilities[v] > probabilities[best]) best = v; | |
| generated += 1; | |
| const eos = token === manifest.config.eos_token_id; | |
| if (!eos) logits = engine.step(token); | |
| const { cloud, raster, active } = sample(); | |
| self.postMessage({ | |
| type: 'token', | |
| token, | |
| eos, | |
| index: generated, | |
| probability, | |
| topProbability: probabilities[best], | |
| meanAbs: engine.meanAbsState, | |
| active, | |
| milliseconds: performance.now() - started, | |
| cloud, | |
| raster, | |
| }, [cloud.buffer, raster.buffer]); | |
| } | |
| self.onmessage = async (event) => { | |
| const message = event.data; | |
| try { | |
| if (message.type === 'load') await load(message); | |
| else if (message.type === 'prime') prime(message.promptIds); | |
| else if (message.type === 'next') next(message); | |
| } catch (error) { | |
| self.postMessage({ type: 'error', message: String(error && error.message ? error.message : error) }); | |
| } | |
| }; | |