Files
wehub-resource-sync 311a3666c4
Build documentation / build (push) Failing after 0s
Unit tests / build (18) (push) Has been cancelled
Unit tests / build (20) (push) Has been cancelled
Unit tests / build (22) (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:44:39 +08:00

540 lines
21 KiB
JavaScript

import { ModelRegistry } from "../../src/transformers.js";
import { getModelFile } from "../../src/utils/hub.js";
import { MAX_TEST_EXECUTION_TIME, DEFAULT_MODEL_OPTIONS } from "../init.js";
const LLAMA_MODEL_ID = "hf-internal-testing/tiny-random-LlamaForCausalLM";
const BERT_MODEL_ID = "hf-internal-testing/tiny-random-BertModel";
const VIT_MODEL_ID = "hf-internal-testing/tiny-random-vit";
// Dedicated model IDs for cache clearing tests to avoid interference with other parallel tests.
// These must NOT be used in any other test file.
const CLEAR_CACHE_MODEL_ID = "onnx-internal-testing/tiny-random-BertModel-ONNX";
const CLEAR_PIPELINE_CACHE_MODEL_ID = "onnx-internal-testing/tiny-random-Qwen3ForCausalLM";
describe("Cache", () => {
describe("ModelRegistry", () => {
describe("get_files", () => {
it(
"should return files for a decoder-only model",
async () => {
const files = await ModelRegistry.get_files(LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(Array.isArray(files)).toBe(true);
expect(files.length).toBeGreaterThan(0);
expect(files).toContain("config.json");
expect(files).toContain("generation_config.json");
expect(files.some((f) => f.startsWith("onnx/") && f.endsWith(".onnx"))).toBe(true);
expect(files).toContain("tokenizer.json");
expect(files).toContain("tokenizer_config.json");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should return files for an encoder-only model",
async () => {
const files = await ModelRegistry.get_files(BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(Array.isArray(files)).toBe(true);
expect(files.length).toBeGreaterThan(0);
expect(files).toContain("config.json");
expect(files.some((f) => f.startsWith("onnx/") && f.endsWith(".onnx"))).toBe(true);
expect(files).toContain("tokenizer.json");
expect(files).toContain("tokenizer_config.json");
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("get_model_files", () => {
it(
"should return model files for a decoder-only model",
async () => {
const files = await ModelRegistry.get_model_files(LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(Array.isArray(files)).toBe(true);
expect(files).toContain("config.json");
expect(files).toContain("generation_config.json");
expect(files.some((f) => f.startsWith("onnx/") && f.endsWith(".onnx"))).toBe(true);
// Should not include tokenizer files
expect(files).not.toContain("tokenizer.json");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should return model files for an encoder-only model",
async () => {
const files = await ModelRegistry.get_model_files(BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(Array.isArray(files)).toBe(true);
expect(files).toContain("config.json");
expect(files.some((f) => f.startsWith("onnx/") && f.endsWith(".onnx"))).toBe(true);
// Encoder-only models should not have generation_config.json
expect(files).not.toContain("generation_config.json");
// Should not include tokenizer files
expect(files).not.toContain("tokenizer.json");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should use model_file_name when provided",
async () => {
const files = await ModelRegistry.get_model_files(BERT_MODEL_ID, {
...DEFAULT_MODEL_OPTIONS,
model_file_name: "custom_model",
});
expect(files).toContain("config.json");
// Should use custom model file name
expect(files.some((f) => f.includes("custom_model") && f.endsWith(".onnx"))).toBe(true);
// Should NOT contain the default 'model' name
expect(files.some((f) => f === "onnx/model.onnx")).toBe(false);
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("get_tokenizer_files", () => {
it(
"should return tokenizer files for a decoder-only model",
async () => {
const files = await ModelRegistry.get_tokenizer_files(LLAMA_MODEL_ID);
expect(files).toEqual(["tokenizer.json", "tokenizer_config.json"]);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should return tokenizer files for an encoder-only model",
async () => {
const files = await ModelRegistry.get_tokenizer_files(BERT_MODEL_ID);
expect(files).toEqual(["tokenizer.json", "tokenizer_config.json"]);
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("get_processor_files", () => {
it(
"should return empty array for text-only models",
async () => {
const llamaFiles = await ModelRegistry.get_processor_files(LLAMA_MODEL_ID);
expect(llamaFiles).toEqual([]);
const bertFiles = await ModelRegistry.get_processor_files(BERT_MODEL_ID);
expect(bertFiles).toEqual([]);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should return processor files for a vision model",
async () => {
const files = await ModelRegistry.get_processor_files(VIT_MODEL_ID);
expect(files).toContain("preprocessor_config.json");
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("get_pipeline_files", () => {
it(
"should return files for text-generation pipeline",
async () => {
const files = await ModelRegistry.get_pipeline_files("text-generation", LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(Array.isArray(files)).toBe(true);
expect(files).toContain("config.json");
expect(files).toContain("generation_config.json");
expect(files.some((f) => f.startsWith("onnx/") && f.endsWith(".onnx"))).toBe(true);
expect(files).toContain("tokenizer.json");
expect(files).toContain("tokenizer_config.json");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should return files for feature-extraction pipeline",
async () => {
const files = await ModelRegistry.get_pipeline_files("feature-extraction", BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(Array.isArray(files)).toBe(true);
expect(files).toContain("config.json");
expect(files.some((f) => f.startsWith("onnx/") && f.endsWith(".onnx"))).toBe(true);
expect(files).toContain("tokenizer.json");
expect(files).toContain("tokenizer_config.json");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should resolve task aliases",
async () => {
const files = await ModelRegistry.get_pipeline_files("sentiment-analysis", BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(Array.isArray(files)).toBe(true);
expect(files).toContain("config.json");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should throw for unsupported pipeline task",
async () => {
await expect(ModelRegistry.get_pipeline_files("invalid-nonexistent-task", BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS)).rejects.toThrow("Unsupported pipeline task");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should exclude vision_encoder for text-generation on multimodal model",
async () => {
const files = await ModelRegistry.get_pipeline_files("text-generation", "onnx-community/gemma-3-4b-it-ONNX", {
device: "webgpu",
dtype: "q4f16",
});
expect(Array.isArray(files)).toBe(true);
// Should NOT include vision_encoder
expect(files.some((f) => f.includes("vision_encoder"))).toBe(false);
// Should include only embed_tokens and decoder_model_merged with q4f16 suffix
expect(files).toContain("onnx/embed_tokens_q4f16.onnx");
expect(files).toContain("onnx/decoder_model_merged_q4f16.onnx");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should exclude vision_encoder for text-generation on ImageTextToText model",
async () => {
const files = await ModelRegistry.get_pipeline_files("text-generation", "onnx-community/Qwen3.5-0.8B-ONNX", {
device: "webgpu",
dtype: "q4f16",
});
expect(Array.isArray(files)).toBe(true);
// Should NOT include vision_encoder
expect(files.some((f) => f.includes("vision_encoder"))).toBe(false);
// Should include embed_tokens and decoder_model_merged
expect(files).toContain("onnx/embed_tokens_q4f16.onnx");
expect(files).toContain("onnx/decoder_model_merged_q4f16.onnx");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should exclude audio_encoder for text-generation on AudioTextToText model (Voxtral)",
async () => {
const files = await ModelRegistry.get_pipeline_files("text-generation", "onnx-community/Voxtral-Mini-3B-2507-ONNX", {
device: "webgpu",
dtype: "q4f16",
});
expect(Array.isArray(files)).toBe(true);
// Should NOT include audio_encoder
expect(files.some((f) => f.includes("audio_encoder"))).toBe(false);
// Should include embed_tokens and decoder_model_merged
expect(files).toContain("onnx/embed_tokens_q4f16.onnx");
expect(files).toContain("onnx/decoder_model_merged_q4f16.onnx");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should exclude audio_encoder for text-generation on AudioTextToText model (Ultravox)",
async () => {
const files = await ModelRegistry.get_pipeline_files("text-generation", "onnx-community/ultravox-v0_5-llama-3_2-1b-ONNX", {
device: "webgpu",
dtype: "q4f16",
});
expect(Array.isArray(files)).toBe(true);
// Should NOT include audio_encoder
expect(files.some((f) => f.includes("audio_encoder"))).toBe(false);
// Should include embed_tokens and decoder_model_merged
expect(files).toContain("onnx/embed_tokens_q4f16.onnx");
expect(files).toContain("onnx/decoder_model_merged_q4f16.onnx");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should exclude vision_encoder and audio_encoder for text-generation on ImageAudioTextToText model",
async () => {
const files = await ModelRegistry.get_pipeline_files("text-generation", "onnx-community/gemma-3n-E2B-it-ONNX", {
device: "webgpu",
dtype: "q4f16",
});
expect(Array.isArray(files)).toBe(true);
// Should NOT include vision_encoder or audio_encoder
expect(files.some((f) => f.includes("vision_encoder"))).toBe(false);
expect(files.some((f) => f.includes("audio_encoder"))).toBe(false);
// Should include embed_tokens and decoder_model_merged
expect(files).toContain("onnx/embed_tokens_q4f16.onnx");
expect(files).toContain("onnx/decoder_model_merged_q4f16.onnx");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should not include processor files for text-only tasks",
async () => {
const files = await ModelRegistry.get_pipeline_files("text-generation", LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS);
// text-generation tasks don't use a processor, so no preprocessor_config.json
expect(files).not.toContain("preprocessor_config.json");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should include processor files for image tasks",
async () => {
const files = await ModelRegistry.get_pipeline_files("image-classification", VIT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(files).toContain("preprocessor_config.json");
// image-classification doesn't use a tokenizer
expect(files).not.toContain("tokenizer.json");
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("is_cached", () => {
it(
"should return a boolean",
async () => {
const cached = await ModelRegistry.is_cached(BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(typeof cached).toBe("boolean");
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("is_cached_files", () => {
it(
"should return cache status with correct shape",
async () => {
const status = await ModelRegistry.is_cached_files(BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(status).toHaveProperty("allCached");
expect(typeof status.allCached).toBe("boolean");
expect(status).toHaveProperty("files");
expect(Array.isArray(status.files)).toBe(true);
expect(status.files.length).toBeGreaterThan(0);
for (const entry of status.files) {
expect(entry).toHaveProperty("file");
expect(typeof entry.file).toBe("string");
expect(entry).toHaveProperty("cached");
expect(typeof entry.cached).toBe("boolean");
}
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should agree with is_cached on allCached",
async () => {
const cached = await ModelRegistry.is_cached(BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
const status = await ModelRegistry.is_cached_files(BERT_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(cached).toBe(status.allCached);
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("is_pipeline_cached", () => {
it(
"should return a boolean for text-generation pipeline",
async () => {
const cached = await ModelRegistry.is_pipeline_cached("text-generation", LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(typeof cached).toBe("boolean");
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("is_pipeline_cached_files", () => {
it(
"should return cache status for text-generation pipeline",
async () => {
const status = await ModelRegistry.is_pipeline_cached_files("text-generation", LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(status).toHaveProperty("allCached");
expect(typeof status.allCached).toBe("boolean");
expect(status).toHaveProperty("files");
expect(Array.isArray(status.files)).toBe(true);
expect(status.files.length).toBeGreaterThan(0);
for (const entry of status.files) {
expect(entry).toHaveProperty("file");
expect(entry).toHaveProperty("cached");
}
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should agree with is_pipeline_cached on allCached",
async () => {
const cached = await ModelRegistry.is_pipeline_cached("text-generation", LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS);
const status = await ModelRegistry.is_pipeline_cached_files("text-generation", LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS);
expect(cached).toBe(status.allCached);
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("get_file_metadata", () => {
it(
"should return metadata for an existing file",
async () => {
const metadata = await ModelRegistry.get_file_metadata(BERT_MODEL_ID, "config.json");
expect(metadata).toHaveProperty("exists", true);
expect(metadata).toHaveProperty("size");
expect(typeof metadata.size).toBe("number");
expect(metadata.size).toBeGreaterThan(0);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should return exists=false for a non-existent file",
async () => {
const metadata = await ModelRegistry.get_file_metadata(BERT_MODEL_ID, "nonexistent_file.bin");
expect(metadata).toHaveProperty("exists", false);
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("clear_cache", () => {
it(
"should clear cached files and report results",
async () => {
// Step 1: Pre-cache config.json by downloading it
await getModelFile(CLEAR_CACHE_MODEL_ID, "config.json", true, {});
// Step 2: Clear the cache
const result = await ModelRegistry.clear_cache(CLEAR_CACHE_MODEL_ID, DEFAULT_MODEL_OPTIONS);
// Step 3: Verify response shape
expect(result).toHaveProperty("filesDeleted");
expect(typeof result.filesDeleted).toBe("number");
expect(result).toHaveProperty("filesCached");
expect(typeof result.filesCached).toBe("number");
expect(result).toHaveProperty("files");
expect(Array.isArray(result.files)).toBe(true);
for (const entry of result.files) {
expect(entry).toHaveProperty("file");
expect(typeof entry.file).toBe("string");
expect(entry).toHaveProperty("deleted");
expect(typeof entry.deleted).toBe("boolean");
expect(entry).toHaveProperty("wasCached");
expect(typeof entry.wasCached).toBe("boolean");
}
// Step 4: config.json should have been cached and deleted
// (it was pre-cached in Step 1)
const configEntry = result.files.find((f) => f.file === "config.json");
expect(configEntry?.wasCached).toBe(true);
expect(configEntry?.deleted).toBe(true);
expect(result.filesDeleted).toBeGreaterThan(0);
// NOTE: We don't re-check is_cached here because it internally calls
// get_model_files() -> AutoConfig.from_pretrained(), which re-downloads
// config.json and re-populates the cache as a side effect.
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("clear_pipeline_cache", () => {
it(
"should clear cached pipeline files and report results",
async () => {
// Step 1: Pre-cache config.json by downloading it
await getModelFile(CLEAR_PIPELINE_CACHE_MODEL_ID, "config.json", true, {});
// Step 2: Clear the pipeline cache
const result = await ModelRegistry.clear_pipeline_cache("text-generation", CLEAR_PIPELINE_CACHE_MODEL_ID, DEFAULT_MODEL_OPTIONS);
// Step 3: Verify response shape
expect(result).toHaveProperty("filesDeleted");
expect(typeof result.filesDeleted).toBe("number");
expect(result).toHaveProperty("filesCached");
expect(typeof result.filesCached).toBe("number");
expect(result).toHaveProperty("files");
expect(Array.isArray(result.files)).toBe(true);
for (const entry of result.files) {
expect(entry).toHaveProperty("file");
expect(entry).toHaveProperty("deleted");
expect(entry).toHaveProperty("wasCached");
}
// Step 4: Should include expected pipeline files
const fileNames = result.files.map((f) => f.file);
expect(fileNames).toContain("config.json");
expect(fileNames).toContain("tokenizer.json");
expect(fileNames.some((f) => f.startsWith("onnx/") && f.endsWith(".onnx"))).toBe(true);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should throw for unsupported pipeline task",
async () => {
await expect(ModelRegistry.clear_pipeline_cache("invalid-nonexistent-task", LLAMA_MODEL_ID, DEFAULT_MODEL_OPTIONS)).rejects.toThrow("Unsupported pipeline task");
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("error handling", () => {
it(
"should throw for empty modelId in is_cached",
async () => {
await expect(ModelRegistry.is_cached("")).rejects.toThrow("modelId is required");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should throw for empty modelId in is_cached_files",
async () => {
await expect(ModelRegistry.is_cached_files("")).rejects.toThrow("modelId is required");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should throw for empty modelId in clear_cache",
async () => {
await expect(ModelRegistry.clear_cache("")).rejects.toThrow("modelId is required");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should throw for empty task in is_pipeline_cached",
async () => {
await expect(ModelRegistry.is_pipeline_cached("", BERT_MODEL_ID)).rejects.toThrow("task is required");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should throw for empty task in is_pipeline_cached_files",
async () => {
await expect(ModelRegistry.is_pipeline_cached_files("", BERT_MODEL_ID)).rejects.toThrow("task is required");
},
MAX_TEST_EXECUTION_TIME,
);
it(
"should throw for empty task in clear_pipeline_cache",
async () => {
await expect(ModelRegistry.clear_pipeline_cache("", BERT_MODEL_ID)).rejects.toThrow("task is required");
},
MAX_TEST_EXECUTION_TIME,
);
});
});
});