wehub-resource-sync
311 lines
13 KiB
JavaScript
311 lines
13 KiB
JavaScript
import { pipeline, LlamaForCausalLM, AutoModelForCausalLM, WhisperForConditionalGeneration, Gemma3ForConditionalGeneration, Gemma3nForConditionalGeneration, VoxtralRealtimeForConditionalGeneration } from "../src/transformers.js";
|
|
|
|
import { init, MAX_MODEL_LOAD_TIME, MAX_MODEL_DISPOSE_TIME, DEFAULT_MODEL_OPTIONS } from "./init.js";
|
|
|
|
// Initialise the testing environment
|
|
init();
|
|
|
|
/**
|
|
* Collects progress events during a loader call and returns them.
|
|
* @param {(cb: Function) => Promise<{ dispose(): Promise<void> }>} loader
|
|
* @returns {Promise<{ events: import('../src/utils/core.js').ProgressInfo[], dispose: () => Promise<void> }>}
|
|
*/
|
|
async function collectEvents(loader) {
|
|
/** @type {import('../src/utils/core.js').ProgressInfo[]} */
|
|
const events = [];
|
|
const result = await loader((info) => events.push(info));
|
|
return { events, dispose: () => result.dispose() };
|
|
}
|
|
|
|
/**
|
|
* Validates progress_total events:
|
|
* 1. loaded is monotonically non-decreasing
|
|
* 2. total is constant across all events
|
|
* 3. final progress value is 100
|
|
* @param {Array<Object>} totalEvents
|
|
*/
|
|
function expectValidTotalEvents(totalEvents) {
|
|
expect(totalEvents.length).toBeGreaterThan(0);
|
|
|
|
for (const event of totalEvents) {
|
|
expect(event).toHaveProperty("status", "progress_total");
|
|
expect(event).toHaveProperty("progress");
|
|
expect(event).toHaveProperty("loaded");
|
|
expect(event).toHaveProperty("total");
|
|
expect(event).toHaveProperty("files");
|
|
expect(typeof event.progress).toBe("number");
|
|
expect(event.progress).toBeGreaterThanOrEqual(0);
|
|
expect(event.progress).toBeLessThanOrEqual(100);
|
|
expect(event.loaded).toBeLessThanOrEqual(event.total);
|
|
}
|
|
|
|
// 1. loaded should be monotonically non-decreasing
|
|
for (let i = 1; i < totalEvents.length; i++) {
|
|
expect(totalEvents[i].loaded).toBeGreaterThanOrEqual(totalEvents[i - 1].loaded);
|
|
}
|
|
|
|
// 2. total should be constant across all events
|
|
const expectedTotal = totalEvents[0].total;
|
|
for (const event of totalEvents) {
|
|
expect(event.total).toBe(expectedTotal);
|
|
}
|
|
|
|
// 3. final progress value should be 100
|
|
expect(totalEvents.at(-1).progress).toBe(100);
|
|
expect(totalEvents.at(-1).loaded).toBe(totalEvents.at(-1).total);
|
|
}
|
|
|
|
/**
|
|
* Validates per-file event lifecycle and structure.
|
|
* @param {Array<Object>} events All collected events.
|
|
* @param {string} model_id Expected model name on events.
|
|
* @param {string[]} expectedFiles File paths that must be present in the files map.
|
|
*/
|
|
function expectValidEventLifecycle(events, model_id, expectedFiles) {
|
|
const totalEvents = events.filter((e) => e.status === "progress_total");
|
|
expectValidTotalEvents(totalEvents);
|
|
|
|
// Exact file count in the final progress_total
|
|
const lastFiles = totalEvents.at(-1).files;
|
|
expect(Object.keys(lastFiles).length).toBe(expectedFiles.length);
|
|
|
|
// All expected files are present and fully loaded
|
|
for (const file of expectedFiles) {
|
|
expect(lastFiles).toHaveProperty([file]);
|
|
expect(lastFiles[file].loaded).toBe(lastFiles[file].total);
|
|
}
|
|
|
|
// Every file emits initiate -> ... -> done lifecycle
|
|
const trackedFiles = new Set(events.filter((e) => e.file).map((e) => e.file));
|
|
for (const file of trackedFiles) {
|
|
const fileEvents = events.filter((e) => e.file === file);
|
|
expect(fileEvents[0].status).toBe("initiate");
|
|
expect(fileEvents.at(-1).status).toBe("done");
|
|
}
|
|
|
|
// All events with a name field should reference the correct model
|
|
for (const event of events) {
|
|
if (event.name) {
|
|
expect(event.name).toBe(model_id);
|
|
}
|
|
}
|
|
|
|
// No double-wrapping: at most one progress_total per progress event
|
|
const progressEvents = events.filter((e) => e.status === "progress");
|
|
expect(totalEvents.length).toBeLessThanOrEqual(progressEvents.length);
|
|
}
|
|
|
|
describe("Progress Callbacks", () => {
|
|
// ---- Llama (decoder-only) ----
|
|
// from_pretrained files: config.json, onnx/model.onnx, generation_config.json
|
|
// pipeline files: + tokenizer.json, tokenizer_config.json
|
|
describe("Llama (decoder-only)", () => {
|
|
const model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM";
|
|
|
|
it(
|
|
"pipeline('text-generation')",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => pipeline("text-generation", model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
expectValidEventLifecycle(events, model_id, ["config.json", "onnx/model.onnx", "generation_config.json", "tokenizer.json", "tokenizer_config.json"]);
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
|
|
it(
|
|
"LlamaForCausalLM.from_pretrained()",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => LlamaForCausalLM.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
expectValidEventLifecycle(events, model_id, ["config.json", "onnx/model.onnx", "generation_config.json"]);
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
|
|
it(
|
|
"AutoModelForCausalLM.from_pretrained()",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => AutoModelForCausalLM.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
expectValidEventLifecycle(events, model_id, ["config.json", "onnx/model.onnx", "generation_config.json"]);
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
});
|
|
|
|
// ---- Whisper (encoder-decoder) ----
|
|
// from_pretrained files: config.json, onnx/encoder_model.onnx, onnx/decoder_model_merged.onnx, generation_config.json
|
|
// pipeline files: + tokenizer.json, tokenizer_config.json, preprocessor_config.json
|
|
describe("Whisper (encoder-decoder)", () => {
|
|
const model_id = "onnx-internal-testing/tiny-random-WhisperForConditionalGeneration";
|
|
|
|
it(
|
|
"pipeline('automatic-speech-recognition')",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => pipeline("automatic-speech-recognition", model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
const expectedFiles = ["config.json", "onnx/encoder_model.onnx", "onnx/decoder_model_merged.onnx", "generation_config.json", "tokenizer.json", "tokenizer_config.json", "preprocessor_config.json"];
|
|
|
|
// Each file should be loaded exactly once: pipeline() must not double-fetch
|
|
// tokenizer.json/tokenizer_config.json/preprocessor_config.json that the
|
|
// tokenizer and processor would otherwise each load independently.
|
|
const initiated = events.filter((e) => e.status === "initiate").map((e) => e.file);
|
|
expect(initiated.sort()).toEqual([...expectedFiles].sort());
|
|
|
|
expectValidEventLifecycle(events, model_id, expectedFiles);
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
|
|
it(
|
|
"WhisperForConditionalGeneration.from_pretrained()",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => WhisperForConditionalGeneration.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
expectValidEventLifecycle(events, model_id, ["config.json", "onnx/encoder_model.onnx", "onnx/decoder_model_merged.onnx", "generation_config.json"]);
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
});
|
|
|
|
// ---- Gemma3 (image-text-to-text) ----
|
|
// from_pretrained files: config.json, onnx/embed_tokens.onnx, onnx/embed_tokens.onnx_data,
|
|
// onnx/decoder_model_merged.onnx, onnx/decoder_model_merged.onnx_data,
|
|
// onnx/vision_encoder.onnx, onnx/vision_encoder.onnx_data, generation_config.json
|
|
describe("Gemma3 (image-text-to-text)", () => {
|
|
const model_id = "onnx-internal-testing/tiny-random-Gemma3ForConditionalGeneration";
|
|
|
|
it(
|
|
"Gemma3ForConditionalGeneration.from_pretrained()",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => Gemma3ForConditionalGeneration.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
expectValidEventLifecycle(events, model_id, ["config.json", "onnx/embed_tokens.onnx", "onnx/embed_tokens.onnx_data", "onnx/decoder_model_merged.onnx", "onnx/decoder_model_merged.onnx_data", "onnx/vision_encoder.onnx", "onnx/vision_encoder.onnx_data", "generation_config.json"]);
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
});
|
|
|
|
// ---- Gemma3n (image-audio-text-to-text) ----
|
|
// from_pretrained files: config.json, onnx/embed_tokens.onnx, onnx/embed_tokens.onnx_data,
|
|
// onnx/decoder_model_merged.onnx, onnx/decoder_model_merged.onnx_data,
|
|
// onnx/audio_encoder.onnx, onnx/audio_encoder.onnx_data,
|
|
// onnx/vision_encoder.onnx, onnx/vision_encoder.onnx_data, generation_config.json
|
|
describe("Gemma3n (image-audio-text-to-text)", () => {
|
|
const model_id = "onnx-internal-testing/tiny-random-Gemma3nForConditionalGeneration";
|
|
|
|
it(
|
|
"Gemma3nForConditionalGeneration.from_pretrained()",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => Gemma3nForConditionalGeneration.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
expectValidEventLifecycle(events, model_id, ["config.json", "onnx/embed_tokens.onnx", "onnx/embed_tokens.onnx_data", "onnx/decoder_model_merged.onnx", "onnx/decoder_model_merged.onnx_data", "onnx/audio_encoder.onnx", "onnx/audio_encoder.onnx_data", "onnx/vision_encoder.onnx", "onnx/vision_encoder.onnx_data", "generation_config.json"]);
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
});
|
|
|
|
// ---- VoxtralRealtime (audio-text-to-text) ----
|
|
// from_pretrained files: config.json, onnx/embed_tokens.onnx, onnx/embed_tokens.onnx_data,
|
|
// onnx/decoder_model_merged.onnx, onnx/decoder_model_merged.onnx_data,
|
|
// onnx/audio_encoder.onnx, onnx/audio_encoder.onnx_data, generation_config.json
|
|
describe("VoxtralRealtime (audio-text-to-text)", () => {
|
|
const model_id = "onnx-internal-testing/tiny-random-VoxtralRealtimeForConditionalGeneration";
|
|
|
|
it(
|
|
"VoxtralRealtimeForConditionalGeneration.from_pretrained()",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => VoxtralRealtimeForConditionalGeneration.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
expectValidEventLifecycle(events, model_id, ["config.json", "onnx/embed_tokens.onnx", "onnx/embed_tokens.onnx_data", "onnx/decoder_model_merged.onnx", "onnx/decoder_model_merged.onnx_data", "onnx/audio_encoder.onnx", "onnx/audio_encoder.onnx_data", "generation_config.json"]);
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
});
|
|
|
|
// ---- Edge cases ----
|
|
describe("Edge cases", () => {
|
|
const model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM";
|
|
|
|
it(
|
|
"no progress_total without progress_callback",
|
|
async () => {
|
|
// When no progress_callback is provided, nothing should throw
|
|
const model = await LlamaForCausalLM.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
|
|
await model.dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
|
|
it(
|
|
"per-file progress events have loaded <= total",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => LlamaForCausalLM.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
const progressEvents = events.filter((e) => e.status === "progress");
|
|
for (const event of progressEvents) {
|
|
expect(event.loaded).toBeLessThanOrEqual(event.total);
|
|
expect(event.loaded).toBeGreaterThanOrEqual(0);
|
|
expect(event.total).toBeGreaterThan(0);
|
|
}
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
|
|
it(
|
|
"per-file progress is monotonically non-decreasing",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => LlamaForCausalLM.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
// Group progress events by file and verify monotonicity within each file
|
|
const progressByFile = {};
|
|
for (const event of events.filter((e) => e.status === "progress")) {
|
|
(progressByFile[event.file] ??= []).push(event.loaded);
|
|
}
|
|
for (const loadedValues of Object.values(progressByFile)) {
|
|
for (let i = 1; i < loadedValues.length; i++) {
|
|
expect(loadedValues[i]).toBeGreaterThanOrEqual(loadedValues[i - 1]);
|
|
}
|
|
}
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
|
|
it(
|
|
"progress_total files map is a deep copy (structuredClone)",
|
|
async () => {
|
|
const { events, dispose } = await collectEvents((cb) => LlamaForCausalLM.from_pretrained(model_id, { ...DEFAULT_MODEL_OPTIONS, progress_callback: cb }));
|
|
|
|
const totalEvents = events.filter((e) => e.status === "progress_total");
|
|
// Each progress_total event should have its own files object (not shared references)
|
|
if (totalEvents.length >= 2) {
|
|
expect(totalEvents[0].files).not.toBe(totalEvents[1].files);
|
|
}
|
|
|
|
await dispose();
|
|
},
|
|
MAX_MODEL_LOAD_TIME + MAX_MODEL_DISPOSE_TIME,
|
|
);
|
|
});
|
|
});
|