src/lib/tokenizer.js

/**
 * Pulpie WebUI - EuroBERT Tokenizer
 *
 * Wraps @huggingface/transformers AutoTokenizer for feyninc/pulpie-orange-small.
 * Handles sequence boundary tokens (BOS, EOS, SEP) and block-level tokenization.
 */

import { AutoTokenizer, PreTrainedTokenizer, env } from '@huggingface/transformers';

export const TOKENIZER_MODEL_ID = 'feyninc/pulpie-orange-small';

export const SPECIAL_TOKENS = {
  BOS: '<|begin_of_text|>',
  EOS: '<|end_of_text|>',
  SEP: '<|sep|>',
};

export const TOKEN_IDS = {
  BOS: 128000,
  EOS: 128001,
  SEP: 128256,
};

let _tokenizerInstance = null;
let _tokenizerLoadingPromise = null;

/**
 * Configure special token IDs from the loaded tokenizer instance.
 * @param {PreTrainedTokenizer|AutoTokenizer} tokenizer
 */
function setupSpecialTokens(tokenizer) {
  try {
    const sepId = tokenizer.model?.tokens_to_ids?.get?.(SPECIAL_TOKENS.SEP);
    if (sepId !== undefined && sepId !== null) {
      TOKEN_IDS.SEP = Number(sepId);
    }
    const bosId = tokenizer.model?.tokens_to_ids?.get?.(SPECIAL_TOKENS.BOS);
    if (bosId !== undefined && bosId !== null) {
      TOKEN_IDS.BOS = Number(bosId);
    }
    const eosId = tokenizer.model?.tokens_to_ids?.get?.(SPECIAL_TOKENS.EOS);
    if (eosId !== undefined && eosId !== null) {
      TOKEN_IDS.EOS = Number(eosId);
    }
  } catch (err) {
    console.warn('[Tokenizer] Could not resolve special token IDs dynamically, using defaults:', err);
  }
}

/**
 * Purge any stale or HTML-poisoned responses from Transformers.js browser cache.
 */
async function clearTransformersCache() {
  if (typeof caches !== 'undefined') {
    try {
      await caches.delete('transformers-cache');
    } catch (err) {
      console.warn('[Tokenizer] Could not purge transformers-cache:', err);
    }
  }
}

/**
 * Load EuroBERT tokenizer with direct local assets priority and remote Hugging Face fallback.
 * @param {function({ status: string, progress?: number }): void} [onStatus]
 * @returns {Promise<PreTrainedTokenizer|AutoTokenizer>}
 */
export async function getTokenizer(onStatus = null) {
  if (_tokenizerInstance) return _tokenizerInstance;
  if (_tokenizerLoadingPromise) return _tokenizerLoadingPromise;

  _tokenizerLoadingPromise = (async () => {
    if (onStatus) onStatus({ status: 'Initializing EuroBERT tokenizer...' });

    // Step 0: Purge stale or corrupted 'transformers-cache' in browser
    await clearTransformersCache();
    env.useBrowserCache = false;

    // Step 1: Direct local load of tokenizer JSON files (bypasses Transformers.js URL/cache quirks)
    try {
      if (onStatus) onStatus({ status: 'Loading tokenizer from local assets...' });
      const [tokJsonResp, tokConfigResp] = await Promise.all([
        fetch('/models/feyninc/pulpie-orange-small/tokenizer.json'),
        fetch('/models/feyninc/pulpie-orange-small/tokenizer_config.json'),
      ]);

      if (tokJsonResp.ok && tokConfigResp.ok) {
        const [tokJsonText, tokConfigText] = await Promise.all([
          tokJsonResp.text(),
          tokConfigResp.text(),
        ]);

        if (!tokJsonText.trim().startsWith('<') && !tokConfigText.trim().startsWith('<')) {
          const tokJson = JSON.parse(tokJsonText);
          const tokConfig = JSON.parse(tokConfigText);
          const tokenizer = new PreTrainedTokenizer(tokJson, tokConfig);
          setupSpecialTokens(tokenizer);
          _tokenizerInstance = tokenizer;
          if (onStatus) onStatus({ status: 'Tokenizer ready (local)', progress: 1 });
          return tokenizer;
        }
      }
      console.warn('[Tokenizer] Local JSON files returned non-JSON responses, falling back...');
    } catch (localErr) {
      console.warn('[Tokenizer] Direct local load failed, trying AutoTokenizer fallback:', localErr.message);
    }

    // Step 2: Hugging Face Remote Fallback
    try {
      if (onStatus) onStatus({ status: 'Loading tokenizer from Hugging Face hub...' });
      env.allowLocalModels = false;
      env.allowRemoteModels = true;

      const tokenizer = await AutoTokenizer.from_pretrained(TOKENIZER_MODEL_ID, {
        progress_callback: (p) => {
          if (onStatus && p?.progress) {
            onStatus({
              status: `Loading tokenizer (${Math.round(p.progress * 100)}%)...`,
              progress: p.progress,
            });
          }
        },
      });

      setupSpecialTokens(tokenizer);
      _tokenizerInstance = tokenizer;
      if (onStatus) onStatus({ status: 'Tokenizer ready', progress: 1 });
      return tokenizer;
    } catch (remoteErr) {
      _tokenizerLoadingPromise = null;
      console.error('[Tokenizer] Failed to load AutoTokenizer:', remoteErr);
      throw new Error(`Tokenizer initialization failed: ${remoteErr.message}`);
    }
  })();

  return _tokenizerLoadingPromise;
}

/**
 * Tokenize a single HTML block string without adding BOS/EOS/SEP tokens.
 * @param {string} blockText
 * @param {AutoTokenizer} tokenizer
 * @returns {number[]} Array of token IDs
 */
export function tokenizeBlock(blockText, tokenizer) {
  if (!tokenizer) throw new Error('Tokenizer not initialized');
  const encoded = tokenizer(blockText, {
    add_special_tokens: false,
    return_tensor: false,
  });
  const inputIds = encoded.input_ids;
  if (Array.isArray(inputIds)) {
    return inputIds.map(Number);
  }
  return Array.from(inputIds).map(Number);
}

/**
 * Tokenize multiple block strings in parallel.
 * @param {string[]} blocks
 * @param {AutoTokenizer} tokenizer
 * @returns {number[][]}
 */
export function tokenizeBlocks(blocks, tokenizer) {
  return blocks.map((block) => tokenizeBlock(block, tokenizer));
}