wehub-resource-sync
110 lines
3.4 KiB
JavaScript
110 lines
3.4 KiB
JavaScript
import { env, LlamaForCausalLM, AutoTokenizer } from "../../src/transformers.js";
|
|
import { init, MAX_TEST_EXECUTION_TIME, DEFAULT_MODEL_OPTIONS } from "../init.js";
|
|
|
|
// Initialise the testing environment
|
|
init();
|
|
|
|
/**
|
|
* A naive custom cache implementation that fetches files directly from the
|
|
* Hugging Face Hub and stores them in an internal (in-memory) map.
|
|
* This satisfies the CacheInterface contract (`match` + `put`).
|
|
*/
|
|
class NaiveFetchCache {
|
|
constructor() {
|
|
/** @type {Map<string, Response>} */
|
|
this.cache = new Map();
|
|
}
|
|
|
|
async match(request) {
|
|
const cached = this.cache.get(request);
|
|
if (cached) {
|
|
return cached.clone();
|
|
}
|
|
|
|
// Not in cache — attempt a fresh fetch from the URL.
|
|
try {
|
|
const response = await fetch(request);
|
|
if (response.ok) {
|
|
this.cache.set(request, response);
|
|
return response.clone();
|
|
}
|
|
} catch {
|
|
// Ignore fetch errors (e.g., invalid URLs like local paths) — treat as cache miss
|
|
}
|
|
return undefined;
|
|
}
|
|
|
|
async put(request, response) {
|
|
if (!this.cache.has(request)) {
|
|
this.cache.set(request, response);
|
|
}
|
|
}
|
|
}
|
|
|
|
describe("Custom cache", () => {
|
|
// Store original env values so we can restore them after tests
|
|
const originalUseCustomCache = env.useCustomCache;
|
|
const originalCustomCache = env.customCache;
|
|
const originalUseBrowserCache = env.useBrowserCache;
|
|
const originalUseFSCache = env.useFSCache;
|
|
const originalAllowLocalModels = env.allowLocalModels;
|
|
|
|
beforeAll(() => {
|
|
// Disable all other caching mechanisms so only customCache is used
|
|
env.useCustomCache = true;
|
|
env.customCache = new NaiveFetchCache();
|
|
env.useBrowserCache = false;
|
|
env.useFSCache = false;
|
|
env.allowLocalModels = false;
|
|
});
|
|
|
|
afterAll(() => {
|
|
// Restore original env values
|
|
env.useCustomCache = originalUseCustomCache;
|
|
env.customCache = originalCustomCache;
|
|
env.useBrowserCache = originalUseBrowserCache;
|
|
env.useFSCache = originalUseFSCache;
|
|
env.allowLocalModels = originalAllowLocalModels;
|
|
});
|
|
|
|
it(
|
|
"should load a model using custom cache (standard)",
|
|
async () => {
|
|
const model_id = "onnx-internal-testing/tiny-random-LlamaForCausalLM-ONNX";
|
|
|
|
const tokenizer = await AutoTokenizer.from_pretrained(model_id);
|
|
const model = await LlamaForCausalLM.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
|
|
|
|
const inputs = await tokenizer("Hello");
|
|
const output = await model(inputs);
|
|
|
|
expect(output.logits).toBeDefined();
|
|
const expected_shape = [...inputs.input_ids.dims, model.config.vocab_size];
|
|
expect(output.logits.dims).toEqual(expected_shape);
|
|
|
|
await model.dispose();
|
|
},
|
|
MAX_TEST_EXECUTION_TIME,
|
|
);
|
|
|
|
it(
|
|
"should load a model using custom cache (external data)",
|
|
async () => {
|
|
const model_id = "onnx-internal-testing/tiny-random-LlamaForCausalLM-ONNX_external";
|
|
|
|
const tokenizer = await AutoTokenizer.from_pretrained(model_id);
|
|
const model = await LlamaForCausalLM.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
|
|
|
|
const inputs = await tokenizer("Hello");
|
|
const output = await model(inputs);
|
|
|
|
expect(output.logits).toBeDefined();
|
|
const expected_shape = [...inputs.input_ids.dims, model.config.vocab_size];
|
|
expect(output.logits.dims).toEqual(expected_shape);
|
|
|
|
await model.dispose();
|
|
},
|
|
MAX_TEST_EXECUTION_TIME,
|
|
);
|
|
});
|