test/model-precision.test.js

import assert from 'node:assert/strict';
import test from 'node:test';
import { selectEmbeddingPrecision } from '../src/model-precision.js';

const manifest = {
  runtime: {
    dtype: { decoder_model_merged: 'q4', vision_encoder: 'q4', embed_tokens: 'fp32' },
    externalDataChunks: { 'decoder_model_merged_q4.onnx': 5, 'vision_encoder_q4.onnx': 1 },
  },
};

test('selects FP16 embeddings and their external shard when shader-f16 is available', () => {
  const selected = selectEmbeddingPrecision(manifest, ['shader-f16']);
  assert.equal(selected.precision, 'fp16');
  assert.equal(selected.manifest.runtime.dtype.embed_tokens, 'fp16');
  assert.equal(selected.manifest.runtime.externalDataChunks['embed_tokens_fp16.onnx'], 1);
  assert.equal(manifest.runtime.dtype.embed_tokens, 'fp32');
});

test('keeps FP32 embeddings when shader-f16 is unavailable', () => {
  const selected = selectEmbeddingPrecision(manifest, new Set(['timestamp-query']));
  assert.equal(selected.precision, 'fp32');
  assert.equal(selected.manifest.runtime.dtype.embed_tokens, 'fp32');
  assert.equal(selected.manifest.runtime.externalDataChunks['embed_tokens_fp16.onnx'], undefined);
});