fix(embedding): resolve ONNX dtype from model declaration or cache

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;

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<string, unknown>): 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 });
+    }
+  });
+});