@flaxia/node 0.3.7 → 0.5.1

package/dist/assets/ai-inference.js

@@ -1,6 +1,238 @@
-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 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: i
+	},
+	"onnx-community/Qwen3-Embedding-0.6B-ONNX": {
+		repo: "onnx-community/Qwen3-Embedding-0.6B-ONNX",
+		revision: "main",
+		maxBytes: i
+	}
+}).map(([e, t]) => [e, { ...t }])), s = [
+	"wasm",
+	"cpu",
+	"webgpu"
+], c = [
+	"q4",
+	"q4f16",
+	"q8",
+	"int8",
+	"uint8"
+], l = {
+	max_new_tokens: {
+		min: 1,
+		max: 1024
+	},
+	temperature: {
+		min: 0,
+		max: 4
+	},
+	top_p: {
+		min: 0,
+		max: 1
+	},
+	top_k: {
+		min: 0,
+		max: 1e3
+	},
+	repetition_penalty: {
+		min: 0,
+		max: 4
+	}
+}, u = 1e5;
+function d() {
+	return [...o.keys()];
+}
+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 = o.get(e);
+	if (!t) throw Error(`ai-inference: model not allowed: ${e} (allowed: ${d().join(", ") || "none"})`);
+	return t;
+}
+function p(e) {
+	if (e != null) {
+		if (typeof e != "string" || !s.includes(e)) throw Error(`ai-inference: device not allowed: ${JSON.stringify(e)} (allowed: ${s.join(", ")})`);
+		return e;
+	}
+}
+function m(e) {
+	if (e != null) {
+		if (typeof e != "string" || !c.includes(e)) throw Error(`ai-inference: dtype not allowed: ${JSON.stringify(e)} (allowed: ${c.join(", ")})`);
+		return e;
+	}
+}
+function h(e) {
+	if (typeof e == "string") {
+		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 ${u} characters`);
+		}
+		return e;
+	}
+	throw Error("ai-inference: input must be a string or an array of strings");
+}
+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)}`);
+			e[n] = i;
+		}
+	}
+	if (t.do_sample !== void 0) {
+		if (typeof t.do_sample != "boolean") throw Error("ai-inference: option do_sample must be a boolean");
+		e.do_sample = t.do_sample;
+	}
+	for (let n of ["src_lang", "tgt_lang"]) {
+		let r = t[n];
+		if (r != null) {
+			if (typeof r != "string" || !/^[a-z]{2,3}(_[A-Za-z0-9]{2,8})?$/.test(r)) throw Error(`ai-inference: option ${n} must be a language code like 'en_XX'`);
+			e[n] = r;
+		}
+	}
+}
+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 r = [
+var y = [
 	"text-classification",
 	"token-classification",
 	"question-answering",
@@ -26,51 +258,115 @@ var r = [
 	"depth-estimation",
 	"feature-extraction",
 	"image-feature-extraction"
-], i = /* @__PURE__ */ new Map(), a = 2, o = () => {
-	i.clear();
-}, s = (e, t) => {
-	if (!t.tokenBuffer) return e;
-	let n = "", r = null, i = t.tokenBufferIntervalMs ?? 50, a = () => {
-		n &&= (e(n), ""), r = null;
+], 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) => {
+			c = !0, s = e;
+		}).finally(() => {
+			a--, o -= t.length;
+		});
+	}, d = () => {
+		r !== null && clearTimeout(r), r = null;
+		let e = n;
+		n = "", e && !c && u(e);
+	}, f = (e) => {
+		if (!l) {
+			if (c) throw s;
+			if (!t.tokenBuffer) {
+				u(e);
+				return;
+			}
+			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));
+		}
+	}, p = async (e) => {
+		if (l = !0, r !== null && clearTimeout(r), r = null, e && !c ? d() : n = "", await i, c) throw s;
 	};
-	return (e) => {
-		n += e, r ||= setTimeout(a, i);
+	return {
+		callback: f,
+		finish: () => p(!0),
+		cancel: () => p(!1)
 	};
 };
-t.backends?.onnx?.wasm && (t.backends.onnx.wasm.numThreads = 1, t.backends.onnx.wasm.proxy = !1);
-var c = async (t, o) => {
-	let { task: c, model: l, input: u, options: d = {} } = t;
-	if (!r.includes(c)) throw Error(`Invalid or unsupported task: ${c}. Supported tasks are: ${r.join(", ")}`);
-	let f = `${c}:${l}`, p = i.get(f);
-	if (!p) {
-		if (i.size >= a) {
-			let e = i.keys().next().value;
-			i.delete(e);
+function w() {
+	let e = globalThis.fetch.bind(globalThis);
+	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 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;
+	};
+}
+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: 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 {
+			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;
 		}
+		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 {
 		let e = {
-			dtype: d.dtype ?? "q4f16",
-			device: d.device
-		}, t = performance.now();
-		console.log(`[flaxia-node] ai-inference: loading pipeline task=${c} model=${l} dtype=${e.dtype} device=${e.device ?? "wasm"}`);
+			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 {
-			p = await n(c, l, e);
-		} catch (t) {
-			if (d.device && d.device !== "wasm") console.warn(`[AiInference] ${d.device} failed, falling back to wasm:`, t), e.device = "wasm", p = await n(c, l, e);
-			else throw console.error(`[flaxia-node] ai-inference: pipeline load FAILED task=${c} model=${l} error=${t instanceof Error ? t.message : String(t)}`), t;
+			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;
 		}
-		console.log(`[flaxia-node] ai-inference: pipeline ready task=${c} model=${l} loadMs=${Math.round(performance.now() - t)}`), i.set(f, p);
+		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();
 	}
-	let m = {
-		max_new_tokens: d.max_new_tokens ?? 128,
-		do_sample: d.do_sample ?? !1
-	};
-	d.temperature != null && (m.temperature = d.temperature), d.top_p != null && (m.top_p = d.top_p), d.top_k != null && (m.top_k = d.top_k), d.repetition_penalty != null && (m.repetition_penalty = d.repetition_penalty), d.src_lang != null && (m.src_lang = d.src_lang), d.tgt_lang != null && (m.tgt_lang = d.tgt_lang), o && p.tokenizer && (m.streamer = new e(p.tokenizer, {
-		skip_prompt: !0,
-		skip_special_tokens: !0,
-		callback_function: s(o, d)
-	}));
-	let h = performance.now(), g = await p(u, m);
-	return console.log(`[flaxia-node] ai-inference: done task=${c} model=${l} execMs=${Math.round(performance.now() - h)} inputLen=${String(u).length}`), { output: g };
 };
 //#endregion
-export { c as handleAiInference, o as releaseCache };
+export { S as getModelCacheSnapshot, T as handleAiInference, x as releaseCache };