@flaxia/node 0.5.0 → 0.5.1

package/dist/assets/ai-inference.js

@@ -1,27 +1,28 @@
-import { TextStreamer as e, env as t, pipeline as n } from "./transformers.web.js";
+import { j as e } from "./webgpuProbe.js";
+import { TextStreamer as t, env as n, pipeline as r } from "./transformers.web.js";
 //#region src/workloads/ai-model-registry.ts
-var r = 512 * 1024 * 1024, i = 4 * 1024 ** 3, a = new Map(Object.entries({
+var i = 512 * 1024 * 1024, a = 4 * 1024 ** 3, o = new Map(Object.entries({
 	"Xenova/bert-base-multilingual-uncased-sentiment": {
 		repo: "Xenova/bert-base-multilingual-uncased-sentiment",
 		revision: "main",
-		maxBytes: r
+		maxBytes: i
 	},
 	"onnx-community/Qwen3-Embedding-0.6B-ONNX": {
 		repo: "onnx-community/Qwen3-Embedding-0.6B-ONNX",
 		revision: "main",
-		maxBytes: r
+		maxBytes: i
 	}
-}).map(([e, t]) => [e, { ...t }])), o = [
+}).map(([e, t]) => [e, { ...t }])), s = [
 	"wasm",
 	"cpu",
 	"webgpu"
-], s = [
+], c = [
 	"q4",
 	"q4f16",
 	"q8",
 	"int8",
 	"uint8"
-], c = {
+], l = {
 	max_new_tokens: {
 		min: 1,
 		max: 1024
@@ -42,45 +43,45 @@ var r = 512 * 1024 * 1024, i = 4 * 1024 ** 3, a = new Map(Object.entries({
 		min: 0,
 		max: 4
 	}
-}, l = 1e5;
-function u() {
-	return [...a.keys()];
+}, u = 1e5;
+function d() {
+	return [...o.keys()];
 }
-function d(e) {
+function f(e) {
 	if (typeof e != "string" || !e) throw Error(`ai-inference: model must be a non-empty string, got ${JSON.stringify(e)}`);
-	let t = a.get(e);
-	if (!t) throw Error(`ai-inference: model not allowed: ${e} (allowed: ${u().join(", ") || "none"})`);
+	let t = o.get(e);
+	if (!t) throw Error(`ai-inference: model not allowed: ${e} (allowed: ${d().join(", ") || "none"})`);
 	return t;
 }
-function f(e) {
+function p(e) {
 	if (e != null) {
-		if (typeof e != "string" || !o.includes(e)) throw Error(`ai-inference: device not allowed: ${JSON.stringify(e)} (allowed: ${o.join(", ")})`);
+		if (typeof e != "string" || !s.includes(e)) throw Error(`ai-inference: device not allowed: ${JSON.stringify(e)} (allowed: ${s.join(", ")})`);
 		return e;
 	}
 }
-function p(e) {
+function m(e) {
 	if (e != null) {
-		if (typeof e != "string" || !s.includes(e)) throw Error(`ai-inference: dtype not allowed: ${JSON.stringify(e)} (allowed: ${s.join(", ")})`);
+		if (typeof e != "string" || !c.includes(e)) throw Error(`ai-inference: dtype not allowed: ${JSON.stringify(e)} (allowed: ${c.join(", ")})`);
 		return e;
 	}
 }
-function m(e) {
+function h(e) {
 	if (typeof e == "string") {
-		if (e.length > 1e5) throw Error(`ai-inference: input exceeds ${l} characters`);
+		if (e.length > 1e5) throw Error(`ai-inference: input exceeds ${u} characters`);
 		return e;
 	}
 	if (Array.isArray(e)) {
 		if (e.length === 0 || e.length > 64) throw Error("ai-inference: input array must hold 1..64 items");
 		for (let t of e) {
 			if (typeof t != "string") throw Error("ai-inference: input array items must be strings");
-			if (t.length > 1e5) throw Error(`ai-inference: input exceeds ${l} characters`);
+			if (t.length > 1e5) throw Error(`ai-inference: input exceeds ${u} characters`);
 		}
 		return e;
 	}
 	throw Error("ai-inference: input must be a string or an array of strings");
 }
-function h(e, t) {
-	for (let [n, r] of Object.entries(c)) {
+function g(e, t) {
+	for (let [n, r] of Object.entries(l)) {
 		let i = t[n];
 		if (i != null) {
 			if (typeof i != "number" || !Number.isFinite(i) || i < r.min || i > r.max) throw Error(`ai-inference: option ${n} must be a number within ${r.min}..${r.max}, got ${JSON.stringify(i)}`);
@@ -99,9 +100,139 @@ function h(e, t) {
 		}
 	}
 }
+var _ = class {
+	constructor(e = {}) {
+		this.options = e, this.entries = /* @__PURE__ */ new Map();
+	}
+	async acquire(t, n, r, i) {
+		let a = this.entries.get(t);
+		if (a?.state === "ready" && a.references === 0 && a.expiresAt <= this.now() && (await this.disposeEntry(a), a = this.entries.get(t)), a?.state === "failed" && (this.entries.delete(t), a = void 0), a || (await this.evictExpired(), a = this.entries.get(t), a)) return this.acquireEntry(a);
+		let o = Math.max(1, Math.floor(this.options.capacity ?? 2));
+		if (this.entries.size >= o && await this.makeRoom(), a = this.entries.get(t), a) return this.acquireEntry(a);
+		let s = this.now(), c = {
+			key: t,
+			descriptor: { ...n },
+			state: "loading",
+			references: 1,
+			updatedAt: s,
+			lastUsedAt: s,
+			expiresAt: s + this.ttlMs(),
+			promise: Promise.resolve(void 0),
+			dispose: i
+		};
+		this.entries.set(t, c), c.promise = (async () => {
+			try {
+				let e = await r();
+				return c.value = e.value, c.descriptor = {
+					...c.descriptor,
+					...e.descriptor
+				}, c.state === "loading" && (c.state = "ready"), c.updatedAt = this.now(), c.expiresAt = this.now() + this.ttlMs(), e.value;
+			} catch (t) {
+				throw c.state = "failed", c.updatedAt = this.now(), c.expiresAt = this.now() + Math.min(this.ttlMs(), e), t;
+			}
+		})();
+		try {
+			let e = await c.promise;
+			return this.createLease(c, e);
+		} catch (e) {
+			throw c.references = Math.max(0, c.references - 1), e;
+		}
+	}
+	async acquireEntry(e) {
+		if (e.state === "failed") throw this.entries.delete(e.key), Error("model cache entry failed before it could be acquired");
+		if (e.state === "evicting") throw Error("model cache entry is being evicted");
+		e.references++, e.lastUsedAt = this.now(), e.expiresAt = this.now() + this.ttlMs();
+		try {
+			let t = await e.promise;
+			return this.createLease(e, t);
+		} catch (t) {
+			throw e.references = Math.max(0, e.references - 1), t;
+		}
+	}
+	getSnapshot(t = this.now()) {
+		return {
+			version: 1,
+			sampledAt: t,
+			expiresAt: t + e,
+			entries: [...this.entries.values()].slice(0, 8).map((e) => ({
+				...e.descriptor,
+				state: e.state,
+				referenceCount: e.references
+			}))
+		};
+	}
+	async releaseAll() {
+		let e = [];
+		for (let t of this.entries.values()) {
+			if (t.state === "failed") {
+				this.entries.delete(t.key);
+				continue;
+			}
+			t.state = "evicting", t.updatedAt = this.now(), t.references === 0 && t.value !== void 0 && e.push(this.disposeEntry(t));
+		}
+		await Promise.all(e);
+	}
+	createLease(e, t) {
+		let n = !1;
+		return {
+			value: t,
+			updateDescriptor: (t) => {
+				n || (e.descriptor = {
+					...e.descriptor,
+					...t
+				}, e.updatedAt = this.now());
+			},
+			invalidate: () => {
+				n || (e.state = "evicting", e.updatedAt = this.now());
+			},
+			release: async () => {
+				n || (n = !0, e.references = Math.max(0, e.references - 1), e.lastUsedAt = this.now(), e.updatedAt = this.now(), e.references === 0 && e.state === "evicting" ? await this.disposeEntry(e) : e.references === 0 && e.state === "ready" && (e.expiresAt = this.now() + this.ttlMs()));
+			}
+		};
+	}
+	async evictExpired() {
+		let e = this.now();
+		for (let t of [...this.entries.values()]) t.state === "failed" && t.expiresAt <= e ? this.entries.delete(t.key) : t.state === "ready" && t.references === 0 && t.expiresAt <= e && await this.disposeEntry(t);
+	}
+	async makeRoom() {
+		let e = Math.max(1, Math.floor(this.options.capacity ?? 2));
+		for (; this.entries.size >= e;) {
+			let e = [...this.entries.values()].filter((e) => e.references === 0 && (e.state === "failed" || e.state === "ready")).sort((e, t) => Number(t.state === "failed") - Number(e.state === "failed") || e.lastUsedAt - t.lastUsedAt)[0];
+			if (!e) throw Error("model cache capacity is occupied by active entries");
+			e.state === "failed" ? this.entries.delete(e.key) : await this.disposeEntry(e);
+		}
+	}
+	async disposeEntry(e) {
+		if (e.references > 0 || e.state === "loading") return;
+		e.state = "evicting", e.updatedAt = this.now();
+		let t = e.value;
+		e.value = void 0;
+		try {
+			t !== void 0 && await e.dispose(t);
+		} finally {
+			this.entries.get(e.key) === e && this.entries.delete(e.key);
+		}
+	}
+	now() {
+		return this.options.now?.() ?? Date.now();
+	}
+	ttlMs() {
+		return Math.max(1, this.options.ttlMs ?? 18e5);
+	}
+};
+function v(e) {
+	return JSON.stringify([
+		e.task,
+		e.modelId,
+		e.revision,
+		e.dtype,
+		e.backend,
+		e.layerRange ?? null
+	]);
+}
 //#endregion
 //#region src/workloads/ai-inference.ts
-var g = [
+var y = [
 	"text-classification",
 	"token-classification",
 	"question-answering",
@@ -127,90 +258,115 @@ var g = [
 	"depth-estimation",
 	"feature-extraction",
 	"image-feature-extraction"
-], _ = /* @__PURE__ */ new Map(), v = 2, y = () => {
-	_.clear();
-}, b = (e, t) => {
-	let n = "", r = null, i = Promise.resolve(), a, o = !1, s = !1, c = (t) => {
-		i = i.then(() => {
-			if (!o) return e(t);
+], b = new _(), x = () => {
+	b.releaseAll();
+}, S = () => b.getSnapshot(), C = (e, t) => {
+	let n = "", r = null, i = Promise.resolve(), a = 0, o = 0, s, c = !1, l = !1, u = (t) => {
+		if (a >= 64 || o + t.length > 262144) throw c = !0, s = /* @__PURE__ */ Error("ai-inference token consumer is too slow; generation stopped at the bounded stream buffer"), s;
+		a++, o += t.length, i = i.then(() => {
+			if (!c) return e(t);
 		}).then(() => void 0, (e) => {
-			o = !0, a = e;
+			c = !0, s = e;
+		}).finally(() => {
+			a--, o -= t.length;
 		});
-	}, l = () => {
+	}, d = () => {
 		r !== null && clearTimeout(r), r = null;
 		let e = n;
-		n = "", e && !o && c(e);
-	}, u = (e) => {
-		if (!s) {
-			if (o) throw a;
+		n = "", e && !c && u(e);
+	}, f = (e) => {
+		if (!l) {
+			if (c) throw s;
 			if (!t.tokenBuffer) {
-				c(e);
+				u(e);
 				return;
 			}
-			n += e, r === null && (r = setTimeout(l, t.tokenBufferIntervalMs ?? 50));
+			if (n.length + e.length > 262144) throw c = !0, s = /* @__PURE__ */ Error("ai-inference token consumer is too slow; generation stopped at the bounded stream buffer"), s;
+			n += e, r === null && (r = setTimeout(() => {
+				try {
+					d();
+				} catch (e) {
+					c = !0, s = e;
+				}
+			}, t.tokenBufferIntervalMs ?? 50));
 		}
-	}, d = async (e) => {
-		if (s = !0, r !== null && clearTimeout(r), r = null, e && !o ? l() : n = "", await i, o) throw a;
+	}, p = async (e) => {
+		if (l = !0, r !== null && clearTimeout(r), r = null, e && !c ? d() : n = "", await i, c) throw s;
 	};
 	return {
-		callback: u,
-		finish: () => d(!0),
-		cancel: () => d(!1)
+		callback: f,
+		finish: () => p(!0),
+		cancel: () => p(!1)
 	};
 };
-function x() {
+function w() {
 	let e = globalThis.fetch.bind(globalThis);
-	t.fetch = async (t, n) => {
+	n.fetch = async (t, n) => {
 		let r = new URL(String(t));
 		if (!/(^|\.)huggingface\.co$/.test(r.hostname) && !/(^|\.)hf\.co$/.test(r.hostname)) throw Error(`ai-inference: model download blocked from ${r.hostname}`);
-		let a = await e(r, n), o = Number(a.headers.get("content-length"));
-		if (Number.isFinite(o) && o > i) throw Error(`ai-inference: model file too large (${o} bytes > ${i})`);
-		return a;
+		let i = await e(r, n), o = Number(i.headers.get("content-length"));
+		if (Number.isFinite(o) && o > a) throw Error(`ai-inference: model file too large (${o} bytes > ${a})`);
+		return i;
 	};
 }
-x(), t.backends?.onnx?.wasm && (t.backends.onnx.wasm.numThreads = 1, t.backends.onnx.wasm.proxy = !1);
-var S = async (t, r) => {
-	let { task: i, model: a, input: o, options: s = {} } = t;
-	if (!g.includes(i)) throw Error(`Invalid or unsupported task: ${i}. Supported tasks are: ${g.join(", ")}`);
-	let c = d(a), l = f(s.device), u = p(s.dtype), y = m(o), x = {};
-	h(x, s);
-	let S = `${i}:${c.repo}:${u ?? "q4f16"}:${l ?? "wasm"}`, C = _.get(S);
-	if (!C) {
-		if (_.size >= v) {
-			let e = _.keys().next().value;
-			_.delete(e);
-		}
+w(), n.backends?.onnx?.wasm && (n.backends.onnx.wasm.numThreads = 1, n.backends.onnx.wasm.proxy = !1);
+var T = async (e, n) => {
+	let { task: i, model: a, input: o, options: s = {} } = e;
+	if (!y.includes(i)) throw Error(`Invalid or unsupported task: ${i}. Supported tasks are: ${y.join(", ")}`);
+	let c = f(a), l = p(s.device), u = m(s.dtype), d = h(o), _ = {};
+	g(_, s);
+	let x = u ?? "q4f16", S = l ?? "wasm", w = v({
+		task: i,
+		modelId: a,
+		revision: c.revision,
+		dtype: x,
+		backend: S
+	}), T = performance.now(), E = await b.acquire(w, {
+		modelId: a,
+		revision: c.revision,
+		dtype: x,
+		backend: S
+	}, async () => {
 		let e = {
-			dtype: u ?? "q4f16",
-			device: l
-		}, t = performance.now();
+			dtype: x,
+			device: l,
+			revision: c.revision
+		};
 		console.log(`[flaxia-node] ai-inference: loading pipeline task=${i} model=${c.repo} revision=${c.revision} dtype=${e.dtype} device=${l ?? "wasm"}`);
+		let t;
 		try {
-			C = await n(i, c.repo, e);
-		} catch (t) {
-			if (l && l !== "wasm") console.warn(`[AiInference] ${l} failed, falling back to wasm:`, t), e.device = "wasm", C = await n(i, c.repo, e);
-			else throw console.error(`[flaxia-node] ai-inference: pipeline load FAILED task=${i} model=${c.repo} error=${t instanceof Error ? t.message : String(t)}`), t;
+			t = await r(i, c.repo, e);
+		} catch (n) {
+			if (l && l !== "wasm") console.warn(`[AiInference] ${l} failed, falling back to wasm:`, n), e.device = "wasm", t = await r(i, c.repo, e);
+			else throw console.error(`[flaxia-node] ai-inference: pipeline load FAILED task=${i} model=${c.repo} error=${n instanceof Error ? n.message : String(n)}`), n;
 		}
-		console.log(`[flaxia-node] ai-inference: pipeline ready task=${i} model=${c.repo} loadMs=${Math.round(performance.now() - t)}`), _.set(S, C);
-	}
-	let w = {
-		max_new_tokens: x.max_new_tokens ?? 128,
-		do_sample: x.do_sample ?? !1
-	};
-	for (let [e, t] of Object.entries(x)) e === "max_new_tokens" || e === "do_sample" || (w[e] = t);
-	let T = r && C.tokenizer ? b(r, s) : null;
-	T && (w.streamer = new e(C.tokenizer, {
-		skip_prompt: !0,
-		skip_special_tokens: !0,
-		callback_function: T.callback
-	}));
-	let E = performance.now(), D;
+		return console.log(`[flaxia-node] ai-inference: pipeline ready task=${i} model=${c.repo} loadMs=${Math.round(performance.now() - T)}`), {
+			value: t,
+			descriptor: { backend: e.device ?? "wasm" }
+		};
+	}, (e) => e.dispose?.()), D = E.value;
 	try {
-		D = await C(y, w);
-	} catch (e) {
-		throw T && await T.cancel(), e;
+		let e = {
+			max_new_tokens: _.max_new_tokens ?? 128,
+			do_sample: _.do_sample ?? !1
+		};
+		for (let [t, n] of Object.entries(_)) t === "max_new_tokens" || t === "do_sample" || (e[t] = n);
+		let r = n && D.tokenizer ? C(n, s) : null;
+		r && (e.streamer = new t(D.tokenizer, {
+			skip_prompt: !0,
+			skip_special_tokens: !0,
+			callback_function: r.callback
+		}));
+		let a = performance.now(), o;
+		try {
+			o = await D(d, e);
+		} catch (e) {
+			throw /out of memory|oom|allocation failed|could not allocate/i.test(e instanceof Error ? e.message : String(e)) && E.invalidate(), r && await r.cancel(), e;
+		}
+		return await r?.finish(), console.log(`[flaxia-node] ai-inference: done task=${i} model=${c.repo} execMs=${Math.round(performance.now() - a)} inputLen=${String(d).length}`), { output: o };
+	} finally {
+		await E.release();
 	}
-	return await T?.finish(), console.log(`[flaxia-node] ai-inference: done task=${i} model=${c.repo} execMs=${Math.round(performance.now() - E)} inputLen=${String(y).length}`), { output: D };
 };
 //#endregion
-export { S as handleAiInference, y as releaseCache };
+export { S as getModelCacheSnapshot, T as handleAiInference, x as releaseCache };