diff --git a/src/services/embedding.ts b/src/services/embedding.ts index 789594cd..f049f051 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,88 @@ 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"], + ["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 `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, cacheRoot?: string): DeclaredDtype { + try { + const parsed = JSON.parse( + readFileSync(join(modelCacheRoot(model, cacheRoot), "config.json"), "utf-8") + ); + const declared = parsed?.["transformers.js_config"]?.dtype; + 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 { 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. + */ +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 = { @@ -164,12 +247,20 @@ export class EmbeddingService { // Local model path const { pipeline } = await ensureTransformersLoaded(); + 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 }); + } + }); +});