From 0e4b5309c781640077c2705627b6c991dc65c9ee Mon Sep 17 00:00:00 2001 From: yihuishou Date: Fri, 9 Oct 2026 12:28:35 +0800 Subject: [PATCH 1/2] fix(embedding): resolve ONNX dtype from model declaration or cache pipeline() without an explicit dtype resolves device=cpu on Node/Bun, and cpu is missing from transformers.js DEFAULT_DEVICE_DTYPE_MAPPING, so it falls back to fp32 and requests onnx/model.onnx. Quantized-only repos do not ship that file, so the local lookup 404s, the remote fallback runs, and offline hosts fail warmup with 'fetch failed'. Resolve the dtype in two tiers before calling pipeline(): first the repo's own transformers.js_config.dtype declaration, then a probe of the cached ONNX variants in descending precision. When neither applies, dtype is omitted and transformers.js keeps its existing defaults. --- src/services/embedding.ts | 45 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/src/services/embedding.ts b/src/services/embedding.ts index 789594cd..234d2dea 100644 --- a/src/services/embedding.ts +++ b/src/services/embedding.ts @@ -1,5 +1,6 @@ import { CONFIG } from "../config.js"; import { log } from "./logger.js"; +import { existsSync, readFileSync } from "node:fs"; import { join } from "node:path"; import { formatOnnxruntimeInitError, @@ -13,6 +14,48 @@ const TIMEOUT_MS = 30000; const GLOBAL_EMBEDDING_KEY = Symbol.for("opencode-mem.embedding.instance"); const MAX_CACHE_SIZE = 100; +// ONNX weight variants in descending precision, mirroring transformers.js +// DEFAULT_DTYPE_SUFFIX_MAPPING (utils/dtypes.js). fp32 has no filename suffix. +const ONNX_DTYPE_VARIANTS = [ + ["fp32", ""], + ["fp16", "_fp16"], + ["q8", "_quantized"], + ["int8", "_int8"], + ["q4", "_q4"], +] as const; + +/** + * Read the dtype the model repo declares under `transformers.js_config.dtype`. + * Repos published for the transformers.js ecosystem carry that block; plain + * PyTorch/HF repos do not, and a top-level `dtype` (e.g. "bfloat16") belongs to + * PyTorch — not to transformers.js. Outranks detectCachedDtype() so an explicit + * declaration is never overridden by whatever happens to sit on disk. An absent + * or malformed config.json returns undefined and falls through to the next tier. + */ +function readDeclaredDtype(model: string): string | undefined { + try { + const parsed = JSON.parse(readFileSync(join(CONFIG.storagePath, ".cache", model, "config.json"), "utf-8")); + const declared = parsed?.["transformers.js_config"]?.dtype; + return typeof declared === "string" && declared.trim() ? declared : undefined; + } catch { + return undefined; + } +} + +/** + * Resolve the dtype whose ONNX weights are already cached, so a quantized-only + * repo is not asked for `model.onnx` (fp32). Returns undefined when nothing is + * cached yet — transformers.js then falls back to its own defaults, which read + * `transformers.js_config.dtype` from config.json and finally the device default. + */ +function detectCachedDtype(model: string): string | undefined { + const onnxDir = join(CONFIG.storagePath, ".cache", model, "onnx"); + const hit = ONNX_DTYPE_VARIANTS.find(([, suffix]) => + existsSync(join(onnxDir, `model${suffix}.onnx`)) + ); + return hit?.[0]; +} + export type EmbeddingTask = "document" | "query"; export type EmbedOptions = { @@ -164,7 +207,9 @@ export class EmbeddingService { // Local model path const { pipeline } = await ensureTransformersLoaded(); + const dtype = readDeclaredDtype(CONFIG.embeddingModel) ?? detectCachedDtype(CONFIG.embeddingModel); this.pipe = await pipeline("feature-extraction", CONFIG.embeddingModel, { + ...(dtype ? { dtype } : {}), progress_callback: progressCallback, }); this.isWarmedUp = true; From eea9a9176ffb57923fe38f9591dd6fe2d5e8420f Mon Sep 17 00:00:00 2001 From: EyJunge1 <149941075+EyJunge1@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:35:27 +0200 Subject: [PATCH 2/2] fix(embedding): harden ONNX dtype resolution for quantized caches Expand the cache probe to the full transformers.js suffix map, skip pipeline dtype overrides when config uses a per-file object map, and cover resolveLocalOnnxDtype with unit tests. --- src/services/embedding.ts | 78 ++++++++++++++++++----- tests/embedding-onnx-dtype.test.ts | 99 ++++++++++++++++++++++++++++++ 2 files changed, 161 insertions(+), 16 deletions(-) create mode 100644 tests/embedding-onnx-dtype.test.ts diff --git a/src/services/embedding.ts b/src/services/embedding.ts index 234d2dea..f049f051 100644 --- a/src/services/embedding.ts +++ b/src/services/embedding.ts @@ -21,41 +21,81 @@ const ONNX_DTYPE_VARIANTS = [ ["fp16", "_fp16"], ["q8", "_quantized"], ["int8", "_int8"], + ["uint8", "_uint8"], ["q4", "_q4"], + ["q4f16", "_q4f16"], + ["bnb4", "_bnb4"], + ["q2", "_q2"], + ["q2f16", "_q2f16"], + ["q1", "_q1"], + ["q1f16", "_q1f16"], ] as const; +/** Concrete dtype strings accepted by transformers.js `pipeline({ dtype })`. */ +export type LocalOnnxDtype = (typeof ONNX_DTYPE_VARIANTS)[number][0] | "auto"; + +type DeclaredDtype = + { kind: "string"; value: LocalOnnxDtype } | { kind: "object" } | { kind: "absent" }; + +function modelCacheRoot(model: string, cacheRoot?: string): string { + return join(cacheRoot ?? join(CONFIG.storagePath, ".cache"), model); +} + /** - * Read the dtype the model repo declares under `transformers.js_config.dtype`. - * Repos published for the transformers.js ecosystem carry that block; plain - * PyTorch/HF repos do not, and a top-level `dtype` (e.g. "bfloat16") belongs to - * PyTorch — not to transformers.js. Outranks detectCachedDtype() so an explicit - * declaration is never overridden by whatever happens to sit on disk. An absent - * or malformed config.json returns undefined and falls through to the next tier. + * Read `transformers.js_config.dtype` from the cached config.json. + * String declarations outrank the cache probe. Object (per-file) maps must not + * be overridden — transformers.js resolves them after loading config. Absent or + * malformed configs fall through to the cache probe. */ -function readDeclaredDtype(model: string): string | undefined { +function readDeclaredDtype(model: string, cacheRoot?: string): DeclaredDtype { try { - const parsed = JSON.parse(readFileSync(join(CONFIG.storagePath, ".cache", model, "config.json"), "utf-8")); + const parsed = JSON.parse( + readFileSync(join(modelCacheRoot(model, cacheRoot), "config.json"), "utf-8") + ); const declared = parsed?.["transformers.js_config"]?.dtype; - return typeof declared === "string" && declared.trim() ? declared : undefined; + if (typeof declared === "string" && declared.trim()) { + return { kind: "string", value: declared.trim() as LocalOnnxDtype }; + } + if (declared && typeof declared === "object" && !Array.isArray(declared)) { + return { kind: "object" }; + } + return { kind: "absent" }; } catch { - return undefined; + return { kind: "absent" }; } } /** * Resolve the dtype whose ONNX weights are already cached, so a quantized-only * repo is not asked for `model.onnx` (fp32). Returns undefined when nothing is - * cached yet — transformers.js then falls back to its own defaults, which read - * `transformers.js_config.dtype` from config.json and finally the device default. + * cached yet — transformers.js then falls back to its own defaults. */ -function detectCachedDtype(model: string): string | undefined { - const onnxDir = join(CONFIG.storagePath, ".cache", model, "onnx"); +function detectCachedDtype(model: string, cacheRoot?: string): LocalOnnxDtype | undefined { + const onnxDir = join(modelCacheRoot(model, cacheRoot), "onnx"); const hit = ONNX_DTYPE_VARIANTS.find(([, suffix]) => existsSync(join(onnxDir, `model${suffix}.onnx`)) ); return hit?.[0]; } +/** + * Resolve a concrete dtype to pass into transformers.js `pipeline()`, or + * undefined to leave dtype selection to transformers.js defaults / config. + * + * @param model Hugging Face model id (cache subdirectory under `.cache`) + * @param cacheRoot Optional override for `{storagePath}/.cache` (tests) + */ +export function resolveLocalOnnxDtype( + model: string, + cacheRoot?: string +): LocalOnnxDtype | undefined { + const declared = readDeclaredDtype(model, cacheRoot); + if (declared.kind === "string") return declared.value; + // Per-file object maps: do not pass options.dtype — it would override config. + if (declared.kind === "object") return undefined; + return detectCachedDtype(model, cacheRoot); +} + export type EmbeddingTask = "document" | "query"; export type EmbedOptions = { @@ -207,14 +247,20 @@ export class EmbeddingService { // Local model path const { pipeline } = await ensureTransformersLoaded(); - const dtype = readDeclaredDtype(CONFIG.embeddingModel) ?? detectCachedDtype(CONFIG.embeddingModel); + const dtype = resolveLocalOnnxDtype(CONFIG.embeddingModel); + if (dtype) { + log("Resolved local ONNX dtype for embedding warmup", { + model: CONFIG.embeddingModel, + dtype, + }); + } this.pipe = await pipeline("feature-extraction", CONFIG.embeddingModel, { ...(dtype ? { dtype } : {}), progress_callback: progressCallback, }); this.isWarmedUp = true; this.initError = null; - log("Embedding model warmed up", { model: CONFIG.embeddingModel }); + log("Embedding model warmed up", { model: CONFIG.embeddingModel, dtype: dtype ?? null }); } catch (error) { const rewritten = formatOnnxruntimeInitError(error); this.initPromise = null; diff --git a/tests/embedding-onnx-dtype.test.ts b/tests/embedding-onnx-dtype.test.ts new file mode 100644 index 00000000..867bbd86 --- /dev/null +++ b/tests/embedding-onnx-dtype.test.ts @@ -0,0 +1,99 @@ +import { describe, expect, it } from "bun:test"; +import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { resolveLocalOnnxDtype } from "../src/services/embedding.js"; + +const MODEL = "org/quantized-embed"; + +function withCache(setup: (cacheRoot: string, modelDir: string, onnxDir: string) => void): string { + const cacheRoot = mkdtempSync(join(tmpdir(), "opencode-mem-onnx-dtype-")); + const modelDir = join(cacheRoot, MODEL); + const onnxDir = join(modelDir, "onnx"); + mkdirSync(onnxDir, { recursive: true }); + setup(cacheRoot, modelDir, onnxDir); + return cacheRoot; +} + +function writeConfig(modelDir: string, transformersJsConfig: Record): void { + writeFileSync( + join(modelDir, "config.json"), + JSON.stringify({ model_type: "bert", "transformers.js_config": transformersJsConfig }), + "utf-8" + ); +} + +function touchOnnx(onnxDir: string, fileName: string): void { + writeFileSync(join(onnxDir, fileName), "stub"); +} + +describe("resolveLocalOnnxDtype", () => { + it("detects q8 from model_quantized.onnx when no config dtype is declared", () => { + const cacheRoot = withCache((_cacheRoot, _modelDir, onnxDir) => { + touchOnnx(onnxDir, "model_quantized.onnx"); + }); + try { + expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBe("q8"); + } finally { + rmSync(cacheRoot, { recursive: true, force: true }); + } + }); + + it("prefers declared string dtype over missing fp32 weights", () => { + const cacheRoot = withCache((_cacheRoot, modelDir, onnxDir) => { + writeConfig(modelDir, { dtype: "q8" }); + touchOnnx(onnxDir, "model_quantized.onnx"); + }); + try { + expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBe("q8"); + } finally { + rmSync(cacheRoot, { recursive: true, force: true }); + } + }); + + it("does not override per-file object dtype with a cache probe", () => { + const cacheRoot = withCache((_cacheRoot, modelDir, onnxDir) => { + writeConfig(modelDir, { dtype: { model: "q8" } }); + touchOnnx(onnxDir, "model_quantized.onnx"); + }); + try { + expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBeUndefined(); + } finally { + rmSync(cacheRoot, { recursive: true, force: true }); + } + }); + + it("prefers fp32 when both fp32 and q8 weights are cached", () => { + const cacheRoot = withCache((_cacheRoot, _modelDir, onnxDir) => { + touchOnnx(onnxDir, "model.onnx"); + touchOnnx(onnxDir, "model_quantized.onnx"); + }); + try { + expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBe("fp32"); + } finally { + rmSync(cacheRoot, { recursive: true, force: true }); + } + }); + + it("returns undefined for an empty cache", () => { + const cacheRoot = withCache(() => { + /* empty model dir with only onnx/ */ + }); + try { + expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBeUndefined(); + } finally { + rmSync(cacheRoot, { recursive: true, force: true }); + } + }); + + it("detects q4f16 from model_q4f16.onnx", () => { + const cacheRoot = withCache((_cacheRoot, _modelDir, onnxDir) => { + touchOnnx(onnxDir, "model_q4f16.onnx"); + }); + try { + expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBe("q4f16"); + } finally { + rmSync(cacheRoot, { recursive: true, force: true }); + } + }); +});