Spaces:
Running
Running
| import type { | |
| WorkerRequest, | |
| WorkerRequestPayload, | |
| WorkerResponse, | |
| WorkerSuccessPayload, | |
| } from "./protocol"; | |
| import type { | |
| LoadProgress, | |
| TokenizationResult, | |
| TokenizerMetadata, | |
| } from "./types"; | |
| interface PendingRequest { | |
| resolve: (payload: WorkerSuccessPayload) => void; | |
| reject: (error: Error) => void; | |
| onProgress?: (progress: LoadProgress) => void; | |
| } | |
| export class TokenizerClient { | |
| private readonly worker = new Worker( | |
| new URL("./tokenizer.worker.ts", import.meta.url), | |
| { type: "module" }, | |
| ); | |
| private readonly pending = new Map<number, PendingRequest>(); | |
| private nextId = 1; | |
| constructor() { | |
| this.worker.onmessage = (event: MessageEvent<WorkerResponse>) => { | |
| const message = event.data; | |
| const request = this.pending.get(message.id); | |
| if (!request) return; | |
| if (message.type === "progress") { | |
| request.onProgress?.(message.progress); | |
| return; | |
| } | |
| this.pending.delete(message.id); | |
| if (message.type === "error") { | |
| request.reject(new Error(message.message)); | |
| } else { | |
| request.resolve(message.payload); | |
| } | |
| }; | |
| this.worker.onerror = (event) => { | |
| const error = new Error(event.message || "Tokenizer Worker 启动失败。"); | |
| this.pending.forEach((request) => request.reject(error)); | |
| this.pending.clear(); | |
| }; | |
| } | |
| private request( | |
| request: WorkerRequestPayload, | |
| onProgress?: (progress: LoadProgress) => void, | |
| ): Promise<WorkerSuccessPayload> { | |
| const id = this.nextId++; | |
| return new Promise((resolve, reject) => { | |
| this.pending.set(id, { resolve, reject, onProgress }); | |
| this.worker.postMessage({ ...request, id } as WorkerRequest); | |
| }); | |
| } | |
| async load( | |
| reference: string, | |
| onProgress: (progress: LoadProgress) => void, | |
| ): Promise<TokenizerMetadata> { | |
| return (await this.request({ type: "load", reference }, onProgress)) as TokenizerMetadata; | |
| } | |
| async tokenize( | |
| text: string, | |
| addSpecialTokens: boolean, | |
| ): Promise<TokenizationResult> { | |
| return (await this.request({ type: "tokenize", text, addSpecialTokens })) as TokenizationResult; | |
| } | |
| async cacheSize(): Promise<number> { | |
| return (await this.request({ type: "cache-size" })) as number; | |
| } | |
| async clearCache(): Promise<void> { | |
| await this.request({ type: "clear-cache" }); | |
| } | |
| dispose(): void { | |
| this.worker.terminate(); | |
| const error = new Error("Tokenizer Worker 已关闭。"); | |
| this.pending.forEach((request) => request.reject(error)); | |
| this.pending.clear(); | |
| } | |
| } | |