From 9c7f27df4691611fc907b68c8108e8c8e530a98f Mon Sep 17 00:00:00 2001 From: "Robert Kent Jr." Date: Tue, 11 Aug 2026 16:42:43 -0400 Subject: [PATCH 1/2] feat(embeddings): make transformers.js embeddings configurable The transformers.js provider pinned `dtype: 'fp32'`, never passed a `device`, and threw for any model outside the dimension registry. Three changes, all backward compatible: Expose `device` and `dtype`. Transformers.js defaults `device` to `cpu` under Node, leaving the accelerated backends that onnxruntime-node already bundles unreachable. `device` is only forwarded when set, and `dtype` keeps its `fp32` default, so device resolution is unchanged when both are omitted. Probe dimensions when the registry has no entry, instead of throwing "Unknown dimensions for model X". The Ollama provider already did this; without it any Hugging Face repo outside the registry was unusable even though the pipeline loaded fine. Registry hits and an explicit `dimensions` option skip the probe. Repoint the `bge-large-en-v1.5` preset from `onnx-community/bge-large-en-v1.5`, whose weights return 401, to `Xenova/bge-large-en-v1.5`, which carries the same weights plus quantized and fp16 variants. `resolveModelForPreset` and `getModelDimensions` both succeeded for the broken preset, so nothing surfaced it until the model failed to load. Benchmarked on an Apple M5 Max, 120 documents, best of 3 after warm-up (docs/sec): model cpu coreml webgpu bge-small-en-v1.5 664 198 1713 bge-base-en-v1.5 198 68 580 Xenova/bge-large-en-v1.5 71 9 201 webgpu is 2.6-2.9x faster than cpu at every size; coreml is 3-8x slower because it falls back to CPU for unsupported ops, and coreml+fp16 fails to load on this onnxruntime build. The ranking is hardware-specific, which is why this is a caller choice rather than a new default. The README documents the measurements so `device` does not get cargo-culted. --- README.md | 35 +++++++++++++ src/embeddings/model-info.ts | 2 +- src/embeddings/transformers-js.ts | 61 +++++++++++++++++++-- test/embeddings-transformers-js.test.ts | 70 +++++++++++++++++++++++++ test/model-info.test.ts | 40 ++++++++++++++ 5 files changed, 204 insertions(+), 4 deletions(-) create mode 100644 test/embeddings-transformers-js.test.ts create mode 100644 test/model-info.test.ts diff --git a/README.md b/README.md index 58a18d7..b9532fd 100644 --- a/README.md +++ b/README.md @@ -324,6 +324,41 @@ ollama({ model: 'nomic-embed-text' }) transformersJs({ model: 'Xenova/all-MiniLM-L6-v2' }) ``` +### Transformers.js device and quantization + +By default Transformers.js runs on the CPU at `fp32`. Pass `device` to offload +inference, and `dtype` to trade accuracy for speed and memory: + +```ts +// Apple Silicon: run on the GPU / Neural Engine +transformersJs({ model: 'bge-base-en-v1.5', device: 'coreml' }) + +// Quantized weights: smaller and faster, slightly less accurate +transformersJs({ model: 'bge-base-en-v1.5', dtype: 'q8' }) +``` + +`device` accepts any Transformers.js device (`auto`, `cpu`, `webgpu`, `coreml`, +`cuda`, `dml`, `webnn`, …) and `dtype` any supported quantization (`fp32`, +`fp16`, `q8`, `q4`, …). Availability depends on the platform and the +`onnxruntime-node` build; unsupported combinations fall back or throw at model +load, so verify on your target before relying on one. + +> **Pick the device by measurement, not by name.** Benchmarked on an Apple +> M5 Max, 120 documents, best of 3 after warm-up (docs/sec, higher is better): +> +> | Model | `cpu` fp32 | `cpu` q8 | `coreml` fp32 | `webgpu` fp32 | +> |---|---:|---:|---:|---:| +> | `bge-small-en-v1.5` | 664 | 652 | 198 | **1713** | +> | `bge-base-en-v1.5` | 198 | 244 | 68 | **580** | +> | `Xenova/bge-large-en-v1.5` | 71 | 84 | 9 | **201** | +> +> `webgpu` was 2.6-2.9x faster than CPU across all three. `coreml` was 3-8x +> *slower* — it falls back to CPU for unsupported ops and pays for the graph +> partitioning, with run-to-run variance up to 10x on the larger model. +> `coreml` + `fp16` fails to load outright on this build (`onnxruntime` graph +> fusion error). Different hardware will rank differently, which is exactly why +> this is a caller choice rather than a default. + ## API ### SearchProvider diff --git a/src/embeddings/model-info.ts b/src/embeddings/model-info.ts index 0667f2e..662a03e 100644 --- a/src/embeddings/model-info.ts +++ b/src/embeddings/model-info.ts @@ -122,7 +122,7 @@ export function getModelMaxTokens(model: string): number | undefined { const MODEL_MAPPINGS: Record> = { 'transformers.js': { 'bge-base-en-v1.5': 'Xenova/bge-base-en-v1.5', - 'bge-large-en-v1.5': 'onnx-community/bge-large-en-v1.5', + 'bge-large-en-v1.5': 'Xenova/bge-large-en-v1.5', 'bge-small-en-v1.5': 'Xenova/bge-small-en-v1.5', 'bge-m3': 'Xenova/bge-m3', 'all-MiniLM-L6-v2': 'Xenova/all-MiniLM-L6-v2', diff --git a/src/embeddings/transformers-js.ts b/src/embeddings/transformers-js.ts index 629d33e..d331364 100644 --- a/src/embeddings/transformers-js.ts +++ b/src/embeddings/transformers-js.ts @@ -12,11 +12,50 @@ export interface TransformersProgressInfo { total?: number } +/** Execution device supported by Transformers.js */ +export type TransformersDevice + = | 'auto' + | 'cpu' + | 'gpu' + | 'wasm' + | 'webgpu' + | 'cuda' + | 'dml' + | 'coreml' + | 'webnn' + | 'webnn-npu' + | 'webnn-gpu' + | (string & {}) + +/** Quantization level supported by Transformers.js */ +export type TransformersDtype + = | 'auto' + | 'fp32' + | 'fp16' + | 'q8' + | 'int8' + | 'uint8' + | 'q4' + | 'bnb4' + | 'q4f16' + | (string & {}) + export interface TransformersEmbeddingOptions { /** Model name (e.g., 'bge-base-en-v1.5' or 'Xenova/bge-base-en-v1.5') */ model?: string /** Embedding dimensions (auto-detected for known models) */ dimensions?: number + /** + * Execution device. Defaults to the Transformers.js default (CPU in Node). + * Set `'coreml'` on Apple Silicon or `'webgpu'` where available to offload + * inference from the CPU. + */ + device?: TransformersDevice + /** + * Quantization level (default: `'fp32'`). Lower precision such as `'q8'` + * reduces model size and speeds up inference at some cost to accuracy. + */ + dtype?: TransformersDtype /** Called with model download progress (initiate → download → progress → done → ready) */ onProgress?: (info: TransformersProgressInfo) => void } @@ -50,6 +89,12 @@ async function clearCorruptedCache(error: unknown, model: string): Promise = { dtype: 'fp32' } + const pipelineOpts: Record = { dtype: options.dtype ?? 'fp32' } + if (options.device) + pipelineOpts.device = options.device if (options.onProgress) pipelineOpts.progress_callback = options.onProgress @@ -73,9 +120,17 @@ export function transformersJs(options: TransformersEmbeddingOptions = {}): Embe throw err }) - const dimensions = options.dimensions ?? getModelDimensions(model) + // Known models resolve from the registry; anything else is probed with a + // single embedding, matching how the Ollama provider handles unknown + // models. Without this, any Hugging Face repo outside the registry is + // unusable even though the pipeline loads fine. + let dimensions = options.dimensions ?? getModelDimensions(model) + if (!dimensions) { + const probe = await extractor(['dimension probe'], { pooling: 'mean', normalize: true }) + dimensions = (probe.data as Float32Array).length + } if (!dimensions) - throw new Error(`Unknown dimensions for model ${model}. Please specify dimensions option.`) + throw new Error(`Could not determine dimensions for model ${model}. Please specify the dimensions option.`) const embedder: EmbeddingProvider = async (texts) => { const output = await extractor(texts, { pooling: 'mean', normalize: true }) diff --git a/test/embeddings-transformers-js.test.ts b/test/embeddings-transformers-js.test.ts new file mode 100644 index 0000000..3bf0dc1 --- /dev/null +++ b/test/embeddings-transformers-js.test.ts @@ -0,0 +1,70 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { pipelineMock } = vi.hoisted(() => ({ pipelineMock: vi.fn() })) + +vi.mock('@huggingface/transformers', () => ({ + env: { cacheDir: undefined }, + pipeline: pipelineMock, +})) + +const { transformersJs } = await import('../src/embeddings/transformers-js') + +describe('transformersJs pipeline options', () => { + beforeEach(() => { + pipelineMock.mockReset() + pipelineMock.mockResolvedValue(async () => ({ data: new Float32Array(0) })) + }) + + it('defaults to fp32 and leaves device unset', async () => { + await transformersJs({ model: 'bge-small-en-v1.5' }).resolve() + + const [task, model, opts] = pipelineMock.mock.calls[0]! + expect(task).toBe('feature-extraction') + expect(model).toBe('Xenova/bge-small-en-v1.5') + expect(opts).toEqual({ dtype: 'fp32' }) + }) + + it('forwards device and dtype to the pipeline', async () => { + await transformersJs({ + model: 'bge-base-en-v1.5', + device: 'coreml', + dtype: 'q8', + }).resolve() + + const [, model, opts] = pipelineMock.mock.calls[0]! + expect(model).toBe('Xenova/bge-base-en-v1.5') + expect(opts).toMatchObject({ device: 'coreml', dtype: 'q8' }) + }) + + it('forwards device without overriding the default dtype', async () => { + await transformersJs({ model: 'bge-small-en-v1.5', device: 'webgpu' }).resolve() + + const [, , opts] = pipelineMock.mock.calls[0]! + expect(opts).toEqual({ dtype: 'fp32', device: 'webgpu' }) + }) + + it('resolves dimensions for the selected model', async () => { + const resolved = await transformersJs({ model: 'bge-large-en-v1.5' }).resolve() + expect(resolved.dimensions).toBe(1024) + }) + + // Models outside the registry were unusable: dimensions could not be looked + // up, so resolve() threw even though the pipeline loaded fine. + it('probes dimensions for a model missing from the registry', async () => { + pipelineMock.mockResolvedValue(async () => ({ data: new Float32Array(384) })) + + const resolved = await transformersJs({ model: 'some-org/unlisted-model' }).resolve() + + expect(resolved.dimensions).toBe(384) + }) + + it('prefers an explicit dimensions option over probing', async () => { + const extractor = vi.fn(async () => ({ data: new Float32Array(384) })) + pipelineMock.mockResolvedValue(extractor) + + const resolved = await transformersJs({ model: 'some-org/unlisted-model', dimensions: 512 }).resolve() + + expect(resolved.dimensions).toBe(512) + expect(extractor).not.toHaveBeenCalled() + }) +}) diff --git a/test/model-info.test.ts b/test/model-info.test.ts new file mode 100644 index 0000000..e764b0c --- /dev/null +++ b/test/model-info.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, it } from 'vitest' +import { DEFAULT_MODELS, getModelDimensions, getModelMaxTokens, resolveModelForPreset } from '../src/embeddings/model-info' + +describe('transformers.js preset mapping', () => { + const presets = [ + ['bge-small-en-v1.5', 'Xenova/bge-small-en-v1.5', 384], + ['bge-base-en-v1.5', 'Xenova/bge-base-en-v1.5', 768], + ['bge-large-en-v1.5', 'Xenova/bge-large-en-v1.5', 1024], + ['bge-m3', 'Xenova/bge-m3', 1024], + ['all-MiniLM-L6-v2', 'Xenova/all-MiniLM-L6-v2', 384], + ] as const + + it.each(presets)('maps %s to a repo with published weights', (preset, repo, dims) => { + expect(resolveModelForPreset(preset, 'transformers.js')).toBe(repo) + expect(getModelDimensions(preset)).toBe(dims) + }) + + // onnx-community/bge-large-en-v1.5 returns 401 — the weights are not public, + // so the preset resolved fine but failed at model load. + it('does not point bge-large at the unavailable onnx-community repo', () => { + expect(resolveModelForPreset('bge-large-en-v1.5', 'transformers.js')) + .not + .toContain('onnx-community') + }) + + it('passes through fully-qualified repo ids untouched', () => { + expect(resolveModelForPreset('Xenova/bge-base-en-v1.5', 'transformers.js')) + .toBe('Xenova/bge-base-en-v1.5') + }) + + it('resolves dimensions and max tokens through repo prefixes', () => { + expect(getModelDimensions('Xenova/bge-large-en-v1.5')).toBe(1024) + expect(getModelMaxTokens('Xenova/bge-large-en-v1.5')).toBe(512) + }) + + it('keeps the transformers.js default consistent with its mapping', () => { + const fallback = DEFAULT_MODELS['transformers.js'] + expect(getModelDimensions(fallback.model)).toBe(fallback.dimensions) + }) +}) From 2f8b20f90bc434d76a5cd15673ac936a63f17848 Mon Sep 17 00:00:00 2001 From: Harlan Wilton Date: Wed, 12 Aug 2026 15:37:24 +1000 Subject: [PATCH 2/2] fix(embeddings): align transformers runtime options --- README.md | 33 +++++-------------------- package.json | 2 +- src/embeddings/transformers-js.ts | 29 ++++++++++------------ test/embeddings-transformers-js.test.ts | 12 +++++++-- test/model-info.test.ts | 8 ------ 5 files changed, 30 insertions(+), 54 deletions(-) diff --git a/README.md b/README.md index b9532fd..f1c4e16 100644 --- a/README.md +++ b/README.md @@ -324,40 +324,19 @@ ollama({ model: 'nomic-embed-text' }) transformersJs({ model: 'Xenova/all-MiniLM-L6-v2' }) ``` -### Transformers.js device and quantization +### Transformers.js runtime options -By default Transformers.js runs on the CPU at `fp32`. Pass `device` to offload -inference, and `dtype` to trade accuracy for speed and memory: +Retriv passes `device` and `dtype` to Transformers.js. Retriv uses `fp32` when +you omit `dtype`. Transformers.js selects the device when you omit `device`. ```ts -// Apple Silicon: run on the GPU / Neural Engine -transformersJs({ model: 'bge-base-en-v1.5', device: 'coreml' }) +transformersJs({ model: 'bge-base-en-v1.5', device: 'webgpu' }) -// Quantized weights: smaller and faster, slightly less accurate transformersJs({ model: 'bge-base-en-v1.5', dtype: 'q8' }) ``` -`device` accepts any Transformers.js device (`auto`, `cpu`, `webgpu`, `coreml`, -`cuda`, `dml`, `webnn`, …) and `dtype` any supported quantization (`fp32`, -`fp16`, `q8`, `q4`, …). Availability depends on the platform and the -`onnxruntime-node` build; unsupported combinations fall back or throw at model -load, so verify on your target before relying on one. - -> **Pick the device by measurement, not by name.** Benchmarked on an Apple -> M5 Max, 120 documents, best of 3 after warm-up (docs/sec, higher is better): -> -> | Model | `cpu` fp32 | `cpu` q8 | `coreml` fp32 | `webgpu` fp32 | -> |---|---:|---:|---:|---:| -> | `bge-small-en-v1.5` | 664 | 652 | 198 | **1713** | -> | `bge-base-en-v1.5` | 198 | 244 | 68 | **580** | -> | `Xenova/bge-large-en-v1.5` | 71 | 84 | 9 | **201** | -> -> `webgpu` was 2.6-2.9x faster than CPU across all three. `coreml` was 3-8x -> *slower* — it falls back to CPU for unsupported ops and pays for the graph -> partitioning, with run-to-run variance up to 10x on the larger model. -> `coreml` + `fp16` fails to load outright on this build (`onnxruntime` graph -> fusion error). Different hardware will rank differently, which is exactly why -> this is a caller choice rather than a default. +Support depends on your Transformers.js version and runtime. Test the selected +combination on the target system. ## API diff --git a/package.json b/package.json index 8f8be99..0f56842 100644 --- a/package.json +++ b/package.json @@ -140,7 +140,7 @@ "@ai-sdk/google": "^3.0.0", "@ai-sdk/mistral": "^3.0.0", "@ai-sdk/openai": "^3.0.0", - "@huggingface/transformers": "^3.0.0", + "@huggingface/transformers": "^3.0.0 || ^4.0.0", "@libsql/client": "^0.14.0 || ^0.15.0 || ^0.16.0 || ^0.17.0", "@upstash/vector": "^1.0.0", "ai": "^4.0.0 || ^5.0.0 || ^6.0.0", diff --git a/src/embeddings/transformers-js.ts b/src/embeddings/transformers-js.ts index d331364..9a3e6ca 100644 --- a/src/embeddings/transformers-js.ts +++ b/src/embeddings/transformers-js.ts @@ -25,7 +25,7 @@ export type TransformersDevice | 'webnn' | 'webnn-npu' | 'webnn-gpu' - | (string & {}) + | 'webnn-cpu' /** Quantization level supported by Transformers.js */ export type TransformersDtype @@ -38,7 +38,10 @@ export type TransformersDtype | 'q4' | 'bnb4' | 'q4f16' - | (string & {}) + | 'q2' + | 'q2f16' + | 'q1' + | 'q1f16' export interface TransformersEmbeddingOptions { /** Model name (e.g., 'bge-base-en-v1.5' or 'Xenova/bge-base-en-v1.5') */ @@ -46,16 +49,14 @@ export interface TransformersEmbeddingOptions { /** Embedding dimensions (auto-detected for known models) */ dimensions?: number /** - * Execution device. Defaults to the Transformers.js default (CPU in Node). - * Set `'coreml'` on Apple Silicon or `'webgpu'` where available to offload - * inference from the CPU. + * Execution device or per-file device map. + * Transformers.js selects the device when this option is omitted. */ - device?: TransformersDevice + device?: TransformersDevice | Record /** - * Quantization level (default: `'fp32'`). Lower precision such as `'q8'` - * reduces model size and speeds up inference at some cost to accuracy. + * Data type or per-file data type map. Defaults to `'fp32'`. */ - dtype?: TransformersDtype + dtype?: TransformersDtype | Record /** Called with model download progress (initiate → download → progress → done → ready) */ onProgress?: (info: TransformersProgressInfo) => void } @@ -90,10 +91,10 @@ async function clearCorruptedCache(error: unknown, model: string): Promise = { dtype: options.dtype ?? 'fp32' } - if (options.device) + if (options.device !== undefined) pipelineOpts.device = options.device if (options.onProgress) pipelineOpts.progress_callback = options.onProgress @@ -120,10 +121,6 @@ export function transformersJs(options: TransformersEmbeddingOptions = {}): Embe throw err }) - // Known models resolve from the registry; anything else is probed with a - // single embedding, matching how the Ollama provider handles unknown - // models. Without this, any Hugging Face repo outside the registry is - // unusable even though the pipeline loads fine. let dimensions = options.dimensions ?? getModelDimensions(model) if (!dimensions) { const probe = await extractor(['dimension probe'], { pooling: 'mean', normalize: true }) diff --git a/test/embeddings-transformers-js.test.ts b/test/embeddings-transformers-js.test.ts index 3bf0dc1..e39c300 100644 --- a/test/embeddings-transformers-js.test.ts +++ b/test/embeddings-transformers-js.test.ts @@ -36,6 +36,16 @@ describe('transformersJs pipeline options', () => { expect(opts).toMatchObject({ device: 'coreml', dtype: 'q8' }) }) + it('forwards per-file device and dtype maps', async () => { + const device = { 'model.onnx': 'webgpu' } as const + const dtype = { 'model.onnx': 'q8' } as const + + await transformersJs({ model: 'bge-base-en-v1.5', device, dtype }).resolve() + + const [, , opts] = pipelineMock.mock.calls[0]! + expect(opts).toMatchObject({ device, dtype }) + }) + it('forwards device without overriding the default dtype', async () => { await transformersJs({ model: 'bge-small-en-v1.5', device: 'webgpu' }).resolve() @@ -48,8 +58,6 @@ describe('transformersJs pipeline options', () => { expect(resolved.dimensions).toBe(1024) }) - // Models outside the registry were unusable: dimensions could not be looked - // up, so resolve() threw even though the pipeline loaded fine. it('probes dimensions for a model missing from the registry', async () => { pipelineMock.mockResolvedValue(async () => ({ data: new Float32Array(384) })) diff --git a/test/model-info.test.ts b/test/model-info.test.ts index e764b0c..eb9ecfc 100644 --- a/test/model-info.test.ts +++ b/test/model-info.test.ts @@ -15,14 +15,6 @@ describe('transformers.js preset mapping', () => { expect(getModelDimensions(preset)).toBe(dims) }) - // onnx-community/bge-large-en-v1.5 returns 401 — the weights are not public, - // so the preset resolved fine but failed at model load. - it('does not point bge-large at the unavailable onnx-community repo', () => { - expect(resolveModelForPreset('bge-large-en-v1.5', 'transformers.js')) - .not - .toContain('onnx-community') - }) - it('passes through fully-qualified repo ids untouched', () => { expect(resolveModelForPreset('Xenova/bge-base-en-v1.5', 'transformers.js')) .toBe('Xenova/bge-base-en-v1.5')