2023-05-02 11:37:39 +08:00
|
|
|
|
|
|
|
|
2023-05-03 08:37:08 +08:00
|
|
|
import { AutoModel, PreTrainedModel } from '../src/transformers.js';
|
|
|
|
import { MAX_TEST_EXECUTION_TIME } from './init.js';
|
2023-05-02 11:37:39 +08:00
|
|
|
|
|
|
|
// 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
|
2023-05-05 06:05:03 +08:00
|
|
|
let model = await AutoModel.from_pretrained('Xenova/t5-small');
|
2023-05-02 11:37:39 +08:00
|
|
|
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);
|
|
|
|
});
|
|
|
|
|
|
|
|
});
|