File size: 728 Bytes
6fa125c 1bb1792 6fa125c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
// @ts-check
import { pipeline } from '@huggingface/transformers';
/**
* @param {{
* modelName: string,
* device: import('@huggingface/transformers').DeviceType,
* onProgress?: import('@huggingface/transformers').ProgressCallback
* }} _
*/
export async function loadModelCore({
modelName,
device,
onProgress
}) {
// Create a text-generation pipeline. Depending on the model this may
// perform downloads of model weights; the library should report progress
// via its own callbacks if available.
const pipe = await pipeline(
'text-generation',
modelName,
{
device,
progress_callback: (progress) => {
if (onProgress) onProgress(progress);
}
});
return pipe;
} |