Spaces:
Running
Running
| // Runtime-agnostic prediction pipeline: `ort` is onnxruntime-web (browser) or | |
| // onnxruntime-node (tests); `getSession(task)` returns an InferenceSession. | |
| import { | |
| prepareFeatures, buildMembers, prepareClassification, prepareRegression, | |
| trimmedMean, softmax, makeRng, isMissing, parseNumber, isNumericColumn, | |
| } from "./kumo.js"; | |
| export const LIMITS = { context: 5000, query: 20000, classes: 10, chunk: 256 }; | |
| export async function predict({ ort, getSession, ctx, qry, target, task, numEstimators, seed = 0, onProgress = () => {} }) { | |
| if (!ctx.columns.includes(target)) throw new Error(`Target "${target}" not found`); | |
| const features = ctx.columns.filter((c) => c !== target); | |
| if (!features.length) throw new Error("Need at least one feature column besides the target"); | |
| const labeled = ctx.rows.filter((r) => !isMissing(r[target])); | |
| const unlabeled = ctx.rows.filter((r) => isMissing(r[target])); | |
| const queryRows = (qry ? qry.rows : []).concat(unlabeled); | |
| if (!queryRows.length) throw new Error("Nothing to predict: add a query CSV or leave the target blank on some rows"); | |
| if (labeled.length < 2) throw new Error("Need at least 2 labeled context rows"); | |
| if (labeled.length > LIMITS.context) throw new Error(`Browser demo limit: ${LIMITS.context} labeled context rows`); | |
| if (queryRows.length > LIMITS.query) throw new Error(`Browser demo limit: ${LIMITS.query} query rows`); | |
| const yVals = labeled.map((r) => r[target]); | |
| if (task === "auto") { | |
| const numeric = isNumericColumn(yVals); | |
| task = numeric && new Set(yVals.map(Number)).size > 10 ? "regression" : "classification"; | |
| } | |
| if (task === "regression" && yVals.some((v) => !Number.isFinite(parseNumber(v)))) throw new Error("Regression needs a numeric target"); | |
| const E = numEstimators; | |
| const rng = makeRng(seed + 1); | |
| // RNG order mirrors sdm: target first, then feature flips, then column shuffle. | |
| const tgt = task === "classification" ? prepareClassification(yVals, rng, E) : prepareRegression(yVals, rng, E); | |
| if (task === "classification") { | |
| if (tgt.K < 2) throw new Error("Classification needs at least 2 distinct target values"); | |
| if (tgt.K > LIMITS.classes) throw new Error(`Browser demo supports up to ${LIMITS.classes} classes (got ${tgt.K})`); | |
| } | |
| onProgress("Preprocessing…", 0); | |
| const prep = prepareFeatures(features, labeled, queryRows); | |
| if (!prep.columns.length) throw new Error("All feature columns are constant in the context rows"); | |
| const members = buildMembers(prep, E, rng); | |
| const session = await getSession(task); | |
| const n = prep.n, nq = queryRows.length, chunk = LIMITS.chunk; | |
| const nChunks = Math.ceil(nq / chunk); | |
| const outW = task === "classification" ? tgt.K : 999; | |
| const perMember = members.map(() => new Float64Array(nq * outW)); | |
| let step = 0; | |
| for (let m = 0; m < E; m++) { | |
| const { x, F, mask } = members[m]; | |
| const yT = task === "classification" | |
| ? new ort.Tensor("int64", tgt.yFor(m), [n]) | |
| : new ort.Tensor("float32", tgt.yFor(m), [n]); | |
| const maskT = new ort.Tensor("bool", mask, [F]); | |
| for (let c = 0; c < nChunks; c++) { | |
| const q0 = c * chunk, q1 = Math.min(nq, q0 + chunk), R = n + (q1 - q0); | |
| const xb = new Float32Array(R * F); | |
| xb.set(x.subarray(0, n * F), 0); | |
| xb.set(x.subarray((n + q0) * F, (n + q1) * F), n * F); | |
| const out = await session.run({ x: new ort.Tensor("float32", xb, [R, F]), y: yT, categorical_mask: maskT }); | |
| const o = out.logits.data, W = out.logits.dims[1]; | |
| for (let i = 0; i < q1 - q0; i++) { | |
| const row = o.subarray(i * W, (i + 1) * W); | |
| const v = task === "classification" ? tgt.unshift(m, row) : tgt.invert(m, row); | |
| perMember[m].set(v.subarray(0, outW), (q0 + i) * outW); | |
| } | |
| step++; | |
| onProgress(`Estimator ${m + 1}/${E} · chunk ${c + 1}/${nChunks}`, step / (E * nChunks)); | |
| } | |
| } | |
| // Post-process | |
| const columns = [], rows = []; | |
| const truthPresent = queryRows.some((r) => target in r && !isMissing(r[target])); | |
| let metrics = ""; | |
| if (task === "classification") { | |
| const names = tgt.classes; | |
| columns.push(`predicted_${target}`, ...(truthPresent ? [`true_${target}`] : []), "confidence", ...features, ...names.map((c) => `p(${c})`)); | |
| let correct = 0, scored = 0; | |
| for (let i = 0; i < nq; i++) { | |
| const avg = Array.from({ length: tgt.K }, (_, k) => members.reduce((s, _m, m) => s + perMember[m][i * tgt.K + k], 0) / E); | |
| const p = softmax(avg); | |
| let best = 0; for (let k = 1; k < p.length; k++) if (p[k] > p[best]) best = k; | |
| const row = { [`predicted_${target}`]: names[best], confidence: +p[best].toFixed(4) }; | |
| if (truthPresent) { | |
| const t = queryRows[i][target] ?? ""; | |
| row[`true_${target}`] = t; | |
| if (!isMissing(t)) { scored++; if (String(t) === names[best] || (Number(t) === Number(names[best]) && t !== "")) correct++; } | |
| } | |
| for (const f of features) row[f] = queryRows[i][f] ?? ""; | |
| names.forEach((c, k) => (row[`p(${c})`] = +p[k].toFixed(4))); | |
| rows.push(row); | |
| } | |
| if (scored) metrics = `accuracy ${(correct / scored).toFixed(3)} on ${scored} labeled query rows`; | |
| } else { | |
| columns.push(`predicted_${target}`, ...(truthPresent ? [`true_${target}`] : []), "p10", "p90", ...features); | |
| const pick = (i, qi) => trimmedMean(members.map((_m, m) => perMember[m][i * 999 + qi])); | |
| let se = 0, ae = 0, cov = 0, scored = 0; const ts = []; | |
| for (let i = 0; i < nq; i++) { | |
| const med = pick(i, 499), p10 = pick(i, 99), p90 = pick(i, 899); | |
| const row = { [`predicted_${target}`]: +med.toPrecision(6), p10: +p10.toPrecision(6), p90: +p90.toPrecision(6) }; | |
| if (truthPresent) { | |
| const t = queryRows[i][target] ?? ""; | |
| row[`true_${target}`] = t; | |
| const tv = parseNumber(t); | |
| if (Number.isFinite(tv)) { scored++; se += (med - tv) ** 2; ae += Math.abs(med - tv); cov += tv >= p10 && tv <= p90; ts.push([tv, med]); } | |
| } | |
| for (const f of features) row[f] = queryRows[i][f] ?? ""; | |
| rows.push(row); | |
| } | |
| if (scored) { | |
| const mt = ts.reduce((a, [t]) => a + t, 0) / scored, sst = ts.reduce((a, [t]) => a + (t - mt) ** 2, 0); | |
| metrics = `RMSE ${Math.sqrt(se / scored).toFixed(3)} · MAE ${(ae / scored).toFixed(3)} · R² ${(1 - se / sst).toFixed(3)} · 80% interval coverage ${(cov / scored).toFixed(2)} on ${scored} rows`; | |
| } | |
| } | |
| return { task, columns, rows, metrics, nContext: n, nQuery: nq, nFeatures: prep.columns.length }; | |
| } | |