Spaces:
Running
Running
| import * as ort from "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/ort.webgpu.bundle.min.mjs"; | |
| import { predict } from "./predict.js"; | |
| const REPO = "https://huggingface.co/mrfakename/Kumo-Tabular-ONNX/resolve/main/onnx"; | |
| ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/"; | |
| ort.env.wasm.numThreads = self.crossOriginIsolated ? Math.min(4, navigator.hardwareConcurrency || 1) : 1; | |
| const sessions = new Map(); | |
| const post = (type, extra) => self.postMessage({ type, ...extra }); | |
| async function fetchModel(url, label) { | |
| let cache = null; | |
| try { cache = await caches.open("kumo-tabular-onnx-v1"); } catch {} | |
| const hit = cache && (await cache.match(url)); | |
| if (hit) return new Uint8Array(await hit.arrayBuffer()); | |
| const res = await fetch(url); | |
| if (!res.ok) throw new Error(`Download failed (${res.status}) for ${label}`); | |
| const total = +res.headers.get("content-length") || 0; | |
| const reader = res.body.getReader(); | |
| const parts = []; let got = 0; | |
| for (;;) { | |
| const { done, value } = await reader.read(); | |
| if (done) break; | |
| parts.push(value); got += value.length; | |
| post("status", { message: `Downloading ${label}… ${(got / 1e6).toFixed(0)}${total ? ` / ${(total / 1e6).toFixed(0)}` : ""} MB`, progress: total ? got / total : null }); | |
| } | |
| const buf = new Uint8Array(got); let o = 0; | |
| for (const p of parts) { buf.set(p, o); o += p.length; } | |
| if (cache) { try { await cache.put(url, new Response(buf)); } catch {} } | |
| return buf; | |
| } | |
| async function getSession(size, task, backend) { | |
| const key = `${size}_${task}_${backend}`; | |
| if (sessions.has(key)) return sessions.get(key); | |
| // keep at most one session alive to bound memory | |
| for (const [k, s] of sessions) { try { await s.release(); } catch {} sessions.delete(k); } | |
| const label = `${size} ${task} model`; | |
| const bytes = await fetchModel(`${REPO}/${size}_${task}.onnx`, label); | |
| post("status", { message: `Initializing ${label} on ${backend.toUpperCase()}…`, progress: null }); | |
| const providers = backend === "webgpu" ? ["webgpu", "wasm"] : ["wasm"]; | |
| const session = await ort.InferenceSession.create(bytes, { executionProviders: providers, graphOptimizationLevel: "all" }); | |
| sessions.set(key, session); | |
| return session; | |
| } | |
| self.onmessage = async ({ data }) => { | |
| if (data.type !== "predict") return; | |
| let backend = data.backend; | |
| if (backend === "webgpu" && !(navigator.gpu && (await navigator.gpu.requestAdapter().catch(() => null)))) backend = "wasm"; | |
| try { | |
| const t0 = performance.now(); | |
| const res = await predict({ | |
| ort, ctx: data.ctx, qry: data.qry, target: data.target, task: data.task, | |
| numEstimators: data.numEstimators, | |
| getSession: (task) => getSession(data.size, task, backend), | |
| onProgress: (message, progress) => post("status", { message, progress }), | |
| }); | |
| const secs = ((performance.now() - t0) / 1000).toFixed(1); | |
| post("result", { | |
| ...res, | |
| summary: `${res.task} · ${data.size} model on ${backend.toUpperCase()} · ${res.nContext} context rows · ${res.nQuery} predicted · ${res.nFeatures} features · ${secs}s`, | |
| }); | |
| } catch (e) { | |
| console.error(e); | |
| post("error", { message: `Error: ${e.message || e}` }); | |
| } | |
| }; | |