Spaces:
Running
Running
| import { pipeline, TextStreamer } from '@huggingface/transformers'; | |
| type Device = 'webgpu'; | |
| type DType = 'q4f16'; | |
| type ChatMessage = { role: 'system' | 'user' | 'assistant'; content: string }; | |
| type WorkerRequest = | |
| | { type: 'load'; modelId: string; revision: string; device: Device; dtype: DType; contextLimit: number } | |
| | { type: 'generate'; requestId: number; messages: ChatMessage[] }; | |
| type ProgressDetails = { | |
| status: string; | |
| file?: string; | |
| progress?: number; | |
| loaded?: number; | |
| total?: number; | |
| }; | |
| type WorkerResponse = | |
| | { type: 'progress'; details: ProgressDetails } | |
| | { type: 'ready'; modelId: string; device: Device; dtype: DType; contextLimit: number } | |
| | { type: 'prompt-size'; requestId: number; tokens: number; limit: number } | |
| | { type: 'token'; requestId: number; text: string } | |
| | { type: 'complete'; requestId: number; text: string; elapsedMs: number } | |
| | { type: 'too-long'; requestId: number; tokens: number; limit: number } | |
| | { type: 'error'; stage: 'load' | 'generate'; message: string; requestId?: number }; | |
| const MAX_NEW_TOKENS = 192; | |
| const workerScope = self as unknown as { | |
| addEventListener: (type: 'message', listener: (event: MessageEvent<WorkerRequest>) => void) => void; | |
| postMessage: (message: WorkerResponse) => void; | |
| }; | |
| async function createGenerator(request: Extract<WorkerRequest, { type: 'load' }>) { | |
| return pipeline('text-generation', request.modelId, { | |
| revision: request.revision, | |
| device: request.device, | |
| dtype: request.dtype, | |
| progress_callback: (progress) => { | |
| workerScope.postMessage({ | |
| type: 'progress', | |
| details: { | |
| status: progress.status, | |
| file: 'file' in progress ? progress.file : undefined, | |
| progress: 'progress' in progress ? progress.progress : undefined, | |
| loaded: 'loaded' in progress ? progress.loaded : undefined, | |
| total: 'total' in progress ? progress.total : undefined, | |
| }, | |
| }); | |
| }, | |
| }); | |
| } | |
| type TextGenerator = Awaited<ReturnType<typeof createGenerator>>; | |
| let generator: TextGenerator | null = null; | |
| let activeModelId: string | null = null; | |
| let activeContextLimit = 8192; | |
| async function handleRequest(request: WorkerRequest) { | |
| if (request.type === 'load') { | |
| try { | |
| activeContextLimit = request.contextLimit; | |
| if (generator && activeModelId !== request.modelId) { | |
| await generator.dispose(); | |
| generator = null; | |
| } | |
| if (!generator) { | |
| generator = await createGenerator(request); | |
| activeModelId = request.modelId; | |
| } | |
| workerScope.postMessage({ | |
| type: 'ready', | |
| modelId: request.modelId, | |
| device: request.device, | |
| dtype: request.dtype, | |
| contextLimit: activeContextLimit, | |
| }); | |
| } catch (error) { | |
| workerScope.postMessage({ | |
| type: 'error', | |
| stage: 'load', | |
| message: error instanceof Error ? error.message : 'The browser could not load the model.', | |
| }); | |
| } | |
| return; | |
| } | |
| if (!generator) { | |
| workerScope.postMessage({ type: 'error', stage: 'generate', requestId: request.requestId, message: 'Load the model before sending a question.' }); | |
| return; | |
| } | |
| try { | |
| const maxPromptTokens = activeContextLimit - MAX_NEW_TOKENS; | |
| const promptTokens = generator.tokenizer.apply_chat_template(request.messages, { | |
| tokenize: true, | |
| return_tensor: false, | |
| return_dict: false, | |
| add_generation_prompt: true, | |
| }).length; | |
| if (promptTokens > maxPromptTokens) { | |
| workerScope.postMessage({ | |
| type: 'too-long', | |
| requestId: request.requestId, | |
| tokens: promptTokens, | |
| limit: maxPromptTokens, | |
| }); | |
| return; | |
| } | |
| workerScope.postMessage({ | |
| type: 'prompt-size', | |
| requestId: request.requestId, | |
| tokens: promptTokens, | |
| limit: maxPromptTokens, | |
| }); | |
| const startedAt = performance.now(); | |
| const output = await generator(request.messages, { | |
| max_new_tokens: MAX_NEW_TOKENS, | |
| do_sample: false, | |
| streamer: new TextStreamer(generator.tokenizer, { | |
| skip_prompt: true, | |
| skip_special_tokens: true, | |
| callback_function: (text) => workerScope.postMessage({ type: 'token', requestId: request.requestId, text }), | |
| }), | |
| }); | |
| const finalMessage = output[0]?.generated_text.at(-1); | |
| const text = typeof finalMessage?.content === 'string' ? finalMessage.content : ''; | |
| workerScope.postMessage({ | |
| type: 'complete', | |
| requestId: request.requestId, | |
| text, | |
| elapsedMs: Math.round(performance.now() - startedAt), | |
| }); | |
| } catch (error) { | |
| workerScope.postMessage({ | |
| type: 'error', | |
| stage: 'generate', | |
| requestId: request.requestId, | |
| message: error instanceof Error ? error.message : 'The browser could not generate an answer.', | |
| }); | |
| } | |
| } | |
| workerScope.addEventListener('message', (event) => { | |
| void handleRequest(event.data); | |
| }); |