Create hub.test.js

This commit is contained in:
Joshua Lochner 2023-05-02 05:37:39 +02:00
parent 595efc631c
commit 8de18bc1d5
1 changed files with 37 additions and 0 deletions

37
tests/hub.test.js Normal file
View File

@ -0,0 +1,37 @@
import { env, AutoModel, PreTrainedModel } from '../src/transformers.js';
// env.allowRemoteModels = false;
const MAX_TEST_EXECUTION_TIME = 60_000; // 60 seconds
// TODO: Set cache folder to a temp directory
describe('Hub', () => {
describe('Loading models', () => {
it('should load a model from the local cache', async () => {
// 1. Local model exists (doesn't matter about status of remote file since local is tried first)
let model = await AutoModel.from_pretrained('t5-small');
expect(model).toBeInstanceOf(PreTrainedModel);
}, MAX_TEST_EXECUTION_TIME);
it('should load a model from the remote cache', async () => {
// 2. Local model doesn't exist, remote file exists
// This tests that fallback functionality is working
let model = await AutoModel.from_pretrained('Xenova/t5-small-web');
expect(model).toBeInstanceOf(PreTrainedModel);
}, MAX_TEST_EXECUTION_TIME);
it('should fail to load a model', async () => {
// 3. Local model doesn't exist, remote file doesn't exist
// This tests that error handling is working.
await expect(
AutoModel.from_pretrained('Xenova/this-model-does-not-exist')
).rejects
.toBeInstanceOf(Error);
}, MAX_TEST_EXECUTION_TIME);
});
});