Skip to content

Commit df22336

Browse files
yihuishouEyJunge1
andauthored
fix(embedding): resolve ONNX dtype from model declaration or cache (#393)
* 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. * 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. --------- Co-authored-by: EyJunge1 <149941075+EyJunge1@users.noreply.github.com>
1 parent 16d5d63 commit df22336

2 files changed

Lines changed: 191 additions & 1 deletion

File tree

‎src/services/embedding.ts‎

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { CONFIG } from "../config.js";
22
import { log } from "./logger.js";
3+
import { existsSync, readFileSync } from "node:fs";
34
import { join } from "node:path";
45
import {
56
formatOnnxruntimeInitError,
@@ -13,6 +14,88 @@ const TIMEOUT_MS = 30000;
1314
const GLOBAL_EMBEDDING_KEY = Symbol.for("opencode-mem.embedding.instance");
1415
const MAX_CACHE_SIZE = 100;
1516

17+
// ONNX weight variants in descending precision, mirroring transformers.js
18+
// DEFAULT_DTYPE_SUFFIX_MAPPING (utils/dtypes.js). fp32 has no filename suffix.
19+
const ONNX_DTYPE_VARIANTS = [
20+
["fp32", ""],
21+
["fp16", "_fp16"],
22+
["q8", "_quantized"],
23+
["int8", "_int8"],
24+
["uint8", "_uint8"],
25+
["q4", "_q4"],
26+
["q4f16", "_q4f16"],
27+
["bnb4", "_bnb4"],
28+
["q2", "_q2"],
29+
["q2f16", "_q2f16"],
30+
["q1", "_q1"],
31+
["q1f16", "_q1f16"],
32+
] as const;
33+
34+
/** Concrete dtype strings accepted by transformers.js `pipeline({ dtype })`. */
35+
export type LocalOnnxDtype = (typeof ONNX_DTYPE_VARIANTS)[number][0] | "auto";
36+
37+
type DeclaredDtype =
38+
{ kind: "string"; value: LocalOnnxDtype } | { kind: "object" } | { kind: "absent" };
39+
40+
function modelCacheRoot(model: string, cacheRoot?: string): string {
41+
return join(cacheRoot ?? join(CONFIG.storagePath, ".cache"), model);
42+
}
43+
44+
/**
45+
* Read `transformers.js_config.dtype` from the cached config.json.
46+
* String declarations outrank the cache probe. Object (per-file) maps must not
47+
* be overridden — transformers.js resolves them after loading config. Absent or
48+
* malformed configs fall through to the cache probe.
49+
*/
50+
function readDeclaredDtype(model: string, cacheRoot?: string): DeclaredDtype {
51+
try {
52+
const parsed = JSON.parse(
53+
readFileSync(join(modelCacheRoot(model, cacheRoot), "config.json"), "utf-8")
54+
);
55+
const declared = parsed?.["transformers.js_config"]?.dtype;
56+
if (typeof declared === "string" && declared.trim()) {
57+
return { kind: "string", value: declared.trim() as LocalOnnxDtype };
58+
}
59+
if (declared && typeof declared === "object" && !Array.isArray(declared)) {
60+
return { kind: "object" };
61+
}
62+
return { kind: "absent" };
63+
} catch {
64+
return { kind: "absent" };
65+
}
66+
}
67+
68+
/**
69+
* Resolve the dtype whose ONNX weights are already cached, so a quantized-only
70+
* repo is not asked for `model.onnx` (fp32). Returns undefined when nothing is
71+
* cached yet — transformers.js then falls back to its own defaults.
72+
*/
73+
function detectCachedDtype(model: string, cacheRoot?: string): LocalOnnxDtype | undefined {
74+
const onnxDir = join(modelCacheRoot(model, cacheRoot), "onnx");
75+
const hit = ONNX_DTYPE_VARIANTS.find(([, suffix]) =>
76+
existsSync(join(onnxDir, `model${suffix}.onnx`))
77+
);
78+
return hit?.[0];
79+
}
80+
81+
/**
82+
* Resolve a concrete dtype to pass into transformers.js `pipeline()`, or
83+
* undefined to leave dtype selection to transformers.js defaults / config.
84+
*
85+
* @param model Hugging Face model id (cache subdirectory under `.cache`)
86+
* @param cacheRoot Optional override for `{storagePath}/.cache` (tests)
87+
*/
88+
export function resolveLocalOnnxDtype(
89+
model: string,
90+
cacheRoot?: string
91+
): LocalOnnxDtype | undefined {
92+
const declared = readDeclaredDtype(model, cacheRoot);
93+
if (declared.kind === "string") return declared.value;
94+
// Per-file object maps: do not pass options.dtype — it would override config.
95+
if (declared.kind === "object") return undefined;
96+
return detectCachedDtype(model, cacheRoot);
97+
}
98+
1699
export type EmbeddingTask = "document" | "query";
17100

18101
export type EmbedOptions = {
@@ -164,12 +247,20 @@ export class EmbeddingService {
164247

165248
// Local model path
166249
const { pipeline } = await ensureTransformersLoaded();
250+
const dtype = resolveLocalOnnxDtype(CONFIG.embeddingModel);
251+
if (dtype) {
252+
log("Resolved local ONNX dtype for embedding warmup", {
253+
model: CONFIG.embeddingModel,
254+
dtype,
255+
});
256+
}
167257
this.pipe = await pipeline("feature-extraction", CONFIG.embeddingModel, {
258+
...(dtype ? { dtype } : {}),
168259
progress_callback: progressCallback,
169260
});
170261
this.isWarmedUp = true;
171262
this.initError = null;
172-
log("Embedding model warmed up", { model: CONFIG.embeddingModel });
263+
log("Embedding model warmed up", { model: CONFIG.embeddingModel, dtype: dtype ?? null });
173264
} catch (error) {
174265
const rewritten = formatOnnxruntimeInitError(error);
175266
this.initPromise = null;

‎tests/embedding-onnx-dtype.test.ts‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
import { describe, expect, it } from "bun:test";
2+
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
3+
import { tmpdir } from "node:os";
4+
import { join } from "node:path";
5+
import { resolveLocalOnnxDtype } from "../src/services/embedding.js";
6+
7+
const MODEL = "org/quantized-embed";
8+
9+
function withCache(setup: (cacheRoot: string, modelDir: string, onnxDir: string) => void): string {
10+
const cacheRoot = mkdtempSync(join(tmpdir(), "opencode-mem-onnx-dtype-"));
11+
const modelDir = join(cacheRoot, MODEL);
12+
const onnxDir = join(modelDir, "onnx");
13+
mkdirSync(onnxDir, { recursive: true });
14+
setup(cacheRoot, modelDir, onnxDir);
15+
return cacheRoot;
16+
}
17+
18+
function writeConfig(modelDir: string, transformersJsConfig: Record<string, unknown>): void {
19+
writeFileSync(
20+
join(modelDir, "config.json"),
21+
JSON.stringify({ model_type: "bert", "transformers.js_config": transformersJsConfig }),
22+
"utf-8"
23+
);
24+
}
25+
26+
function touchOnnx(onnxDir: string, fileName: string): void {
27+
writeFileSync(join(onnxDir, fileName), "stub");
28+
}
29+
30+
describe("resolveLocalOnnxDtype", () => {
31+
it("detects q8 from model_quantized.onnx when no config dtype is declared", () => {
32+
const cacheRoot = withCache((_cacheRoot, _modelDir, onnxDir) => {
33+
touchOnnx(onnxDir, "model_quantized.onnx");
34+
});
35+
try {
36+
expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBe("q8");
37+
} finally {
38+
rmSync(cacheRoot, { recursive: true, force: true });
39+
}
40+
});
41+
42+
it("prefers declared string dtype over missing fp32 weights", () => {
43+
const cacheRoot = withCache((_cacheRoot, modelDir, onnxDir) => {
44+
writeConfig(modelDir, { dtype: "q8" });
45+
touchOnnx(onnxDir, "model_quantized.onnx");
46+
});
47+
try {
48+
expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBe("q8");
49+
} finally {
50+
rmSync(cacheRoot, { recursive: true, force: true });
51+
}
52+
});
53+
54+
it("does not override per-file object dtype with a cache probe", () => {
55+
const cacheRoot = withCache((_cacheRoot, modelDir, onnxDir) => {
56+
writeConfig(modelDir, { dtype: { model: "q8" } });
57+
touchOnnx(onnxDir, "model_quantized.onnx");
58+
});
59+
try {
60+
expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBeUndefined();
61+
} finally {
62+
rmSync(cacheRoot, { recursive: true, force: true });
63+
}
64+
});
65+
66+
it("prefers fp32 when both fp32 and q8 weights are cached", () => {
67+
const cacheRoot = withCache((_cacheRoot, _modelDir, onnxDir) => {
68+
touchOnnx(onnxDir, "model.onnx");
69+
touchOnnx(onnxDir, "model_quantized.onnx");
70+
});
71+
try {
72+
expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBe("fp32");
73+
} finally {
74+
rmSync(cacheRoot, { recursive: true, force: true });
75+
}
76+
});
77+
78+
it("returns undefined for an empty cache", () => {
79+
const cacheRoot = withCache(() => {
80+
/* empty model dir with only onnx/ */
81+
});
82+
try {
83+
expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBeUndefined();
84+
} finally {
85+
rmSync(cacheRoot, { recursive: true, force: true });
86+
}
87+
});
88+
89+
it("detects q4f16 from model_q4f16.onnx", () => {
90+
const cacheRoot = withCache((_cacheRoot, _modelDir, onnxDir) => {
91+
touchOnnx(onnxDir, "model_q4f16.onnx");
92+
});
93+
try {
94+
expect(resolveLocalOnnxDtype(MODEL, cacheRoot)).toBe("q4f16");
95+
} finally {
96+
rmSync(cacheRoot, { recursive: true, force: true });
97+
}
98+
});
99+
});

0 commit comments

Comments
 (0)