src/tokenizer/tokenizer.worker.ts

import {
  clearTokenizerCache,
  fetchOptionalJson,
  fetchTokenizerAsset,
  tokenizerCacheSize,
} from "./cache";
import { WebTokenizerEngine } from "./engine";
import { refineTokenizerFormat } from "./format-detection";
import {
  buildHubFileUrl,
  findCompanionFile,
  resolveTokenizer,
} from "./hub";
import {
  collectSpecialTokenTexts,
  enrichTemplatesFromConfig,
  extractSpecialTokens,
  extractVocabSize,
  tokenizerType,
} from "./metadata";
import type { WorkerRequest, WorkerResponse } from "./protocol";
import type {
  HubFile,
  LoadProgress,
  RemoteFile,
  SpecialTokenMap,
  TokenizerMetadata,
} from "./types";

const MAX_INPUT_CHARACTERS = 200_000;
const workerScope = self as unknown as {
  onmessage: ((event: MessageEvent<WorkerRequest>) => void) | null;
  postMessage: (message: WorkerResponse) => void;
};

interface ActiveTokenizer {
  engine: WebTokenizerEngine;
  metadata: TokenizerMetadata;
  specialTokens: SpecialTokenMap;
  specialTokenTexts: Set<string>;
  templates: ReturnType<typeof enrichTemplatesFromConfig>;
}

let activeTokenizer: ActiveTokenizer | null = null;

function post(message: WorkerResponse): void {
  workerScope.postMessage(message);
}

function postProgress(id: number, progress: LoadProgress): void {
  post({ id, type: "progress", progress });
}

function friendlyError(error: unknown): string {
  if (error instanceof Error) return error.message;
  return "发生了无法识别的 Tokenizer 错误。";
}

function remoteFile(
  file: HubFile,
  repoId: string,
  commitSha: string,
): RemoteFile {
  return {
    ...file,
    url: buildHubFileUrl(repoId, commitSha, file.filename),
  };
}

function companionFile(
  files: HubFile[],
  primaryFilename: string,
  target: string,
  repoId: string,
  commitSha: string,
): RemoteFile | null {
  const file = findCompanionFile(files, primaryFilename, target);
  return file ? remoteFile(file, repoId, commitSha) : null;
}

async function loadTokenizer(id: number, reference: string): Promise<TokenizerMetadata> {
  activeTokenizer?.engine.dispose();
  activeTokenizer = null;
  const startedAt = performance.now();

  postProgress(id, {
    phase: "resolving",
    loaded: 0,
    total: null,
    message: "正在解析 Hugging Face 仓库",
  });
  const resolved = await resolveTokenizer(reference);
  const source = remoteFile(
    resolved.file,
    resolved.repoId,
    resolved.commitSha,
  );

  const tokenizerConfigFile = companionFile(
    resolved.availableFiles,
    resolved.file.filename,
    "tokenizer_config.json",
    resolved.repoId,
    resolved.commitSha,
  );
  const specialTokensFile = companionFile(
    resolved.availableFiles,
    resolved.file.filename,
    "special_tokens_map.json",
    resolved.repoId,
    resolved.commitSha,
  );
  const modelConfigFile = companionFile(
    resolved.availableFiles,
    resolved.file.filename,
    "config.json",
    resolved.repoId,
    resolved.commitSha,
  );

  const [bytes, downloadedTokenizerConfig, specialTokenConfig, modelConfig] =
    await Promise.all([
      fetchTokenizerAsset(source, (progress) => postProgress(id, progress)),
      fetchOptionalJson(tokenizerConfigFile),
      fetchOptionalJson(specialTokensFile),
      fetchOptionalJson(modelConfigFile),
    ]);

  const format = refineTokenizerFormat(
    resolved.format,
    resolved.file.filename,
    bytes,
  );
  postProgress(id, {
    phase: "initializing",
    loaded: bytes.byteLength,
    total: bytes.byteLength,
    message: `正在初始化 ${format}`,
  });

  const engine = await WebTokenizerEngine.create(bytes, format);
  const tokenizerConfig = {
    ...(resolved.tokenizerConfig ?? {}),
    ...(downloadedTokenizerConfig ?? {}),
  };
  const specialTokens = extractSpecialTokens(
    tokenizerConfig,
    specialTokenConfig,
  );
  const templates = enrichTemplatesFromConfig(engine.templates, tokenizerConfig);
  const metadata: TokenizerMetadata = {
    repoId: resolved.repoId,
    requestedRevision: resolved.requestedRevision,
    commitSha: resolved.commitSha,
    format,
    filename: resolved.file.filename,
    tokenizerType: tokenizerType(format, tokenizerConfig),
    vocabSize: extractVocabSize(modelConfig),
    specialTokens,
    fileSize: bytes.byteLength,
    loadTimeMs: performance.now() - startedAt,
    sourceUrl: source.url,
  };

  activeTokenizer = {
    engine,
    metadata,
    specialTokens,
    specialTokenTexts: collectSpecialTokenTexts(specialTokens),
    templates,
  };
  return metadata;
}

function visibleToken(text: string): string {
  if (!text) return "∅";
  return text
    .replaceAll(" ", "␣")
    .replaceAll("\n", "↵")
    .replaceAll("\r", "␍")
    .replaceAll("\t", "⇥");
}

function tokenize(text: string, addSpecialTokens: boolean) {
  if (!activeTokenizer) throw new Error("请先加载一个 Tokenizer。");
  if (!text) throw new Error("请输入需要 Tokenize 的文本。");
  if (text.length > MAX_INPUT_CHARACTERS) {
    throw new Error("输入文本不能超过 200,000 个字符。");
  }

  const startedAt = performance.now();
  const encoded = activeTokenizer.engine.encode(
    text,
    addSpecialTokens,
    activeTokenizer.templates,
  );
  const tokens = encoded.ids.map((id, index) => {
    const decoded = activeTokenizer?.engine.decodeToken(id) ?? {
      text: "",
      byteDisplay: null,
    };
    return {
      index,
      id,
      text: decoded.text,
      displayText: decoded.byteDisplay ?? visibleToken(decoded.text),
      special:
        encoded.templateSpecialIds.has(id) ||
        activeTokenizer?.specialTokenTexts.has(decoded.text) === true,
    };
  });

  return {
    tokens,
    tokenCount: tokens.length,
    characterCount: text.length,
    elapsedMs: performance.now() - startedAt,
  };
}

workerScope.onmessage = (event) => {
  const request = event.data;
  void (async () => {
    try {
      switch (request.type) {
        case "load":
          post({
            id: request.id,
            type: "success",
            payload: await loadTokenizer(request.id, request.reference),
          });
          break;
        case "tokenize":
          post({
            id: request.id,
            type: "success",
            payload: tokenize(request.text, request.addSpecialTokens),
          });
          break;
        case "cache-size":
          post({
            id: request.id,
            type: "success",
            payload: await tokenizerCacheSize(),
          });
          break;
        case "clear-cache":
          await clearTokenizerCache();
          post({ id: request.id, type: "success", payload: null });
          break;
      }
    } catch (error) {
      post({ id: request.id, type: "error", message: friendlyError(error) });
    }
  })();
};