src/tokenizer/engine.ts

import initializeKitoken, { Kitoken } from "kitoken/web";

import { applySpecialTemplates, type TokenizerTemplate } from "./templates";
import type { TokenizerFormat } from "./types";

let initialization: Promise<unknown> | null = null;

export async function initializeEngine(): Promise<void> {
  initialization ??= initializeKitoken();
  await initialization;
}

function createKitoken(bytes: Uint8Array, format: TokenizerFormat): Kitoken {
  switch (format) {
    case "tokenizers":
      return Kitoken.from_tokenizers(bytes);
    case "sentencepiece":
      return Kitoken.from_sentencepiece(bytes);
    case "tiktoken":
      return Kitoken.from_tiktoken(bytes);
    case "tekken":
      return Kitoken.from_tekken(bytes);
  }
}

function readTemplates(instance: Kitoken): TokenizerTemplate[] {
  const config = instance.config() as unknown;
  if (!config || typeof config !== "object" || !("templates" in config)) return [];
  const templates = (config as { templates?: unknown }).templates;
  if (!Array.isArray(templates)) return [];

  return templates.filter(
    (template): template is TokenizerTemplate =>
      Boolean(
        template &&
          typeof template === "object" &&
          "content" in template &&
          typeof template.content === "string" &&
          "position" in template &&
          typeof template.position === "string",
      ),
  );
}

export class WebTokenizerEngine {
  readonly templates: TokenizerTemplate[];
  private readonly decoder = new TextDecoder();
  private readonly fatalDecoder = new TextDecoder("utf-8", { fatal: true });

  private constructor(private readonly instance: Kitoken) {
    this.templates = readTemplates(instance);
  }

  static async create(
    bytes: Uint8Array,
    format: TokenizerFormat,
  ): Promise<WebTokenizerEngine> {
    await initializeEngine();
    return new WebTokenizerEngine(createKitoken(bytes, format));
  }

  encode(
    text: string,
    addSpecialTokens: boolean,
    templates: TokenizerTemplate[],
  ): { ids: number[]; templateSpecialIds: Set<number> } {
    const ids = Array.from(this.instance.encode(text, true));
    if (!addSpecialTokens) return { ids, templateSpecialIds: new Set() };

    const applied = applySpecialTemplates(ids, templates, (content) =>
      Array.from(this.instance.encode(content, true)),
    );
    return { ids: applied.ids, templateSpecialIds: applied.specialIds };
  }

  decodeToken(id: number): { text: string; byteDisplay: string | null } {
    const bytes = this.instance.decode(Uint32Array.of(id), true);
    const text = this.decoder.decode(bytes);
    try {
      this.fatalDecoder.decode(bytes);
      return { text, byteDisplay: null };
    } catch {
      return {
        text,
        byteDisplay: Array.from(bytes, (byte) =>
          `\\x${byte.toString(16).padStart(2, "0").toUpperCase()}`,
        ).join(""),
      };
    }
  }

  dispose(): void {
    this.instance.free();
  }
}