chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,310 @@
|
||||
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,
|
||||
);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user