El tokenizador en el navegador

// Fixed-vocabulary UCI tokenizer. The vocabulary is a deterministic enumeration (it does not
// depend on any data) implemented identically in Python (`rukh/src/rukh/tokenize/uci_vocab.py`)
// and here; `tests/parity.test.ts` checks both against the fixture exported from Python.

const FILES = 'abcdefgh';
const RANKS = '12345678';

/** The 64 squares in file-major order: a1, a2, ..., a8, b1, ..., h8 (index `file * 8 + rank`). */
export const SQUARES: readonly string[] = Array.from({ length: 64 }, (_, i) => {
  const file = FILES[Math.floor(i / 8)];
  const rank = RANKS[i % 8];
  return `${file}${rank}`;
});

/** Special tokens, in id order (0-7). Elo bins follow (8-61) and moves start at id 62. */
export const SPECIALS: readonly string[] = [
  '<pad>',
  '<bos>',
  '<eos>',
  '<mask>',
  '<unk>',
  '<1-0>',
  '<0-1>',
  '<1/2>',
];

export const PAD_ID = 0;
export const BOS_ID = 1;
export const EOS_ID = 2;
export const MASK_ID = 3;
export const UNK_ID = 4;

export const ELO_MIN = 600;
/** Upper clamp before binning, so 3200+ lands in `<w3200>`; it is not the highest bin (3200). */
export const ELO_CLAMP_MAX = 3299;
export const ELO_BIN_WIDTH = 100;
/** Number of Elo bins per side: 600, 700, ..., 3200. */
export const ELO_BIN_COUNT = 27;

const PROMOTION_PIECES = ['q', 'r', 'b', 'n'] as const;

/** Result strings accepted by `encodeGame`, mapped to their special token. */
const RESULT_TOKENS: Record<string, string> = {
  '1-0': '<1-0>',
  '0-1': '<0-1>',
  '1/2-1/2': '<1/2>',
};

function coords(square: string): [number, number] {
  return [FILES.indexOf(square[0] ?? ''), RANKS.indexOf(square[1] ?? '')];
}

/** True if `to` is reachable from `from` by a queen (same file, rank or diagonal) or a knight. */
function reachable(from: string, to: string): boolean {
  if (from === to) return false;
  const [ff, fr] = coords(from);
  const [tf, tr] = coords(to);
  const df = Math.abs(tf - ff);
  const dr = Math.abs(tr - fr);
  const queen = df === 0 || dr === 0 || df === dr;
  const knight = (df === 1 && dr === 2) || (df === 2 && dr === 1);
  return queen || knight;
}

/** Maps an Elo to its bin token: `clamp(elo, 600, 3299)` floored to hundreds, four digits. */
export function eloBin(elo: number, side: 'w' | 'b'): string {
  const clamped = Math.min(Math.max(elo, ELO_MIN), ELO_CLAMP_MAX);
  const bin = Math.floor(clamped / ELO_BIN_WIDTH) * ELO_BIN_WIDTH;
  return `<${side}${String(bin).padStart(4, '0')}>`;
}

function eloBinTokens(side: 'w' | 'b'): string[] {
  return Array.from({ length: ELO_BIN_COUNT }, (_, i) => eloBin(ELO_MIN + i * ELO_BIN_WIDTH, side));
}

/** Enumerates the 1968 pseudo-legal UCI move strings: 1792 plain moves then 176 promotions. */
export function buildMoves(): string[] {
  const moves: string[] = [];
  for (const from of SQUARES) {
    for (const to of SQUARES) {
      if (reachable(from, to)) moves.push(`${from}${to}`);
    }
  }
  for (const [fromRank, toRank] of [
    ['7', '8'],
    ['2', '1'],
  ]) {
    for (const from of SQUARES) {
      if (from[1] !== fromRank) continue;
      for (const to of SQUARES) {
        if (to[1] !== toRank) continue;
        if (Math.abs(coords(to)[0] - coords(from)[0]) > 1) continue;
        for (const piece of PROMOTION_PIECES) moves.push(`${from}${to}${piece}`);
      }
    }
  }
  return moves;
}

/** The full 2030-token vocabulary: specials, white Elo bins, black Elo bins, moves. */
export function buildVocab(): string[] {
  return [...SPECIALS, ...eloBinTokens('w'), ...eloBinTokens('b'), ...buildMoves()];
}

export class UciTokenizer {
  readonly tokens: readonly string[];
  private readonly ids: Map<string, number>;

  constructor(tokens: readonly string[] = buildVocab()) {
    this.tokens = tokens;
    this.ids = new Map(tokens.map((token, id) => [token, id]));
  }

  get size(): number {
    return this.tokens.length;
  }

  /** Id of a token, or `<unk>` (4) if it is not in the vocabulary. */
  tokenId(token: string): number {
    return this.ids.get(token) ?? UNK_ID;
  }

  /** Encodes a whitespace-separated UCI move string; unknown moves map to `<unk>`. */
  encodeMoves(uci: string): number[] {
    return uci
      .trim()
      .split(/\s+/)
      .filter((move) => move.length > 0)
      .map((move) => this.tokenId(move));
  }

  /**
   * `[<bos>, <wXXXX>, <bXXXX>, ...moves, <result>, <eos>]`. When the whole game fits in `maxLen`
   * the sequence ends with `<eos>`; otherwise it is cut to the first `maxLen` ids (no result, no
   * `<eos>`), exactly as `UciTokenizer.encode_game` does in Python.
   */
  encodeGame(
    uci: string,
    whiteElo: number,
    blackElo: number,
    result: string,
    maxLen = 200,
  ): number[] {
    const resultToken = RESULT_TOKENS[result.trim()];
    if (resultToken === undefined) throw new Error(`unknown result: ${result}`);
    const ids = [
      BOS_ID,
      this.tokenId(eloBin(whiteElo, 'w')),
      this.tokenId(eloBin(blackElo, 'b')),
      ...this.encodeMoves(uci),
      this.tokenId(resultToken),
      EOS_ID,
    ];
    return ids.length <= maxLen ? ids : ids.slice(0, maxLen);
  }

  /** Token strings for a sequence of ids; out-of-range ids decode to `<unk>`. */
  decode(ids: readonly number[]): string[] {
    return ids.map((id) => this.tokens[id] ?? '<unk>');
  }
}

// Character-level tokenizer for numbered SAN text plus the result (`1.e4 e5 2.Nf3 ... 1-0`).
// Mirrors `rukh/src/rukh/tokenize/san_chars.py`: `<pad>`, `<bos>`, `<eos>` first, then the fixed
// 32-symbol alphabet in this exact order; `encode` wraps the text in `<bos>` ... `<eos>`.

export const SAN_ALPHABET = ' #+-.0123456789=BKNOQRabcdefghx/';
export const SAN_SPECIALS: readonly string[] = ['<pad>', '<bos>', '<eos>'];
export const SAN_PAD_ID = 0;
export const SAN_BOS_ID = 1;
export const SAN_EOS_ID = 2;

export class SanCharTokenizer {
  readonly tokens: readonly string[];
  private readonly ids: Map<string, number>;

  constructor(alphabet: string = SAN_ALPHABET, specials: readonly string[] = SAN_SPECIALS) {
    this.tokens = [...specials, ...alphabet];
    this.ids = new Map(this.tokens.map((token, id) => [token, id]));
  }

  get size(): number {
    return this.tokens.length;
  }

  /** `[<bos>, one id per character, <eos>]`; characters outside the alphabet throw. */
  encode(text: string): number[] {
    const ids = Array.from(text, (char) => {
      const id = this.ids.get(char);
      if (id === undefined)
        throw new Error(`character not in SAN alphabet: ${JSON.stringify(char)}`);
      return id;
    });
    return [SAN_BOS_ID, ...ids, SAN_EOS_ID];
  }

  /** Joins the characters back into text; special tokens are dropped. */
  decode(ids: readonly number[]): string {
    return ids
      .map((id) => this.tokens[id] ?? '')
      .filter((token) => !SAN_SPECIALS.includes(token))
      .join('');
  }
}

// Applies a BPE model trained with Hugging Face `tokenizers` (`models.BPE` with a
// `WhitespaceSplit` pre-tokenizer, no continuing-subword prefix and no end-of-word suffix) to
// UCI text. It reads the `tokenizer.json` that `Tokenizer.save()` writes and reproduces the
// encoder: split on whitespace, start from single characters and merge the adjacent pair with
// the lowest merge rank (leftmost on ties) until nothing merges, then map symbols to ids.
// Only that subset is implemented: a file with dropout, `fuse_unk`, `byte_fallback`, an affix or
// another pre-tokenizer is rejected by the constructor instead of being encoded incorrectly.

export interface BpeAddedToken {
  id: number;
  content: string;
  special?: boolean;
}

export interface BpeFile {
  model: {
    type?: string;
    unk_token?: string | null;
    ignore_merges?: boolean;
    dropout?: number | null;
    continuing_subword_prefix?: string | null;
    end_of_word_suffix?: string | null;
    fuse_unk?: boolean;
    byte_fallback?: boolean;
    vocab: Record<string, number>;
    merges: string[] | [string, string][];
  };
  pre_tokenizer?: { type?: string } | null;
  added_tokens?: BpeAddedToken[];
}

export const DEFAULT_UNK_TOKEN = '<unk>';

/** Pre-tokenizers this loader reproduces; anything else splits words differently. */
const SUPPORTED_PRE_TOKENIZERS = ['WhitespaceSplit'];

/**
 * Rejects a `tokenizer.json` whose options this loader does not implement, instead of encoding
 * it as if they held their default value. The supported subset is listed in the README.
 */
function assertSupported(file: BpeFile): void {
  const { model } = file;
  const unsupported = (option: string, value: unknown, expected: string): Error =>
    new Error(
      `unsupported BPE model: ${option} is ${JSON.stringify(value) ?? 'undefined'}, ` +
        `this loader only supports ${expected}`,
    );

  if (model.type !== undefined && model.type !== 'BPE') {
    throw unsupported('model.type', model.type, '"BPE"');
  }
  if (model.dropout !== undefined && model.dropout !== null) {
    throw unsupported('model.dropout', model.dropout, 'null');
  }
  for (const option of ['continuing_subword_prefix', 'end_of_word_suffix'] as const) {
    const value = model[option];
    if (value !== undefined && value !== null && value !== '') {
      throw unsupported(`model.${option}`, value, 'null or ""');
    }
  }
  for (const option of ['fuse_unk', 'byte_fallback'] as const) {
    if (model[option] === true) throw unsupported(`model.${option}`, true, 'false');
  }
  const preTokenizer = file.pre_tokenizer;
  if (preTokenizer != null && !SUPPORTED_PRE_TOKENIZERS.includes(preTokenizer.type ?? '')) {
    throw unsupported(
      'pre_tokenizer.type',
      preTokenizer.type,
      SUPPORTED_PRE_TOKENIZERS.map((type) => `"${type}"`).join(', '),
    );
  }
}

function escapeRegExp(text: string): string {
  return text.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
}

export class BpeTokenizer {
  readonly vocab: ReadonlyMap<string, number>;
  readonly tokens: readonly string[];
  readonly unkId: number;
  /** Merge rank by `a\u0000b` key (the pair separator cannot appear in a token). */
  private readonly ranks: Map<string, number>;
  private readonly addedTokens: Map<string, number>;
  private readonly addedPattern: RegExp | null;
  private readonly ignoreMerges: boolean;

  constructor(file: BpeFile) {
    assertSupported(file);
    const { model } = file;
    this.vocab = new Map(Object.entries(model.vocab));
    const size = Math.max(-1, ...this.vocab.values()) + 1;
    const tokens: string[] = new Array<string>(size).fill('');
    for (const [token, id] of this.vocab) tokens[id] = token;
    this.tokens = tokens;

    const unkToken = model.unk_token ?? DEFAULT_UNK_TOKEN;
    const unkId = this.vocab.get(unkToken);
    if (unkId === undefined) throw new Error(`BPE vocabulary has no ${unkToken} token`);
    this.unkId = unkId;
    this.ignoreMerges = model.ignore_merges ?? false;

    this.ranks = new Map();
    model.merges.forEach((merge, rank) => {
      const [a, b] = typeof merge === 'string' ? splitMerge(merge) : merge;
      if (!this.vocab.has(a) || !this.vocab.has(b) || !this.vocab.has(a + b)) return;
      this.ranks.set(pairKey(a, b), rank);
    });

    this.addedTokens = new Map();
    for (const added of file.added_tokens ?? []) this.addedTokens.set(added.content, added.id);
    const contents = [...this.addedTokens.keys()].sort((x, y) => y.length - x.length);
    this.addedPattern =
      contents.length > 0 ? new RegExp(`(${contents.map(escapeRegExp).join('|')})`) : null;
  }

  get size(): number {
    return this.tokens.length;
  }

  /** Ids of the tokens of `text`, applying added tokens first and BPE to the rest. */
  encode(text: string): number[] {
    return this.tokenize(text).map(
      (token) => this.addedTokens.get(token) ?? this.vocab.get(token) ?? this.unkId,
    );
  }

  /** Token strings of `text` (unknown characters stay as themselves). */
  tokenize(text: string): string[] {
    if (this.addedPattern === null) return this.tokenizeSpan(text);
    const out: string[] = [];
    for (const span of text.split(this.addedPattern)) {
      if (span.length === 0) continue;
      if (this.addedTokens.has(span)) out.push(span);
      else out.push(...this.tokenizeSpan(span));
    }
    return out;
  }

  /** Token strings for a sequence of ids; out-of-range ids decode to the unknown token. */
  decode(ids: readonly number[]): string[] {
    return ids.map((id) => this.tokens[id] || this.tokens[this.unkId] || DEFAULT_UNK_TOKEN);
  }

  private tokenizeSpan(text: string): string[] {
    const out: string[] = [];
    for (const word of text.split(/\s+/)) {
      if (word.length > 0) out.push(...this.mergeWord(word));
    }
    return out;
  }

  private mergeWord(word: string): string[] {
    if (this.ignoreMerges && this.vocab.has(word)) return [word];
    const symbols = Array.from(word);
    for (;;) {
      let best = -1;
      let bestRank = Number.POSITIVE_INFINITY;
      for (let i = 0; i + 1 < symbols.length; i++) {
        const rank = this.ranks.get(pairKey(symbols[i]!, symbols[i + 1]!));
        if (rank !== undefined && rank < bestRank) {
          bestRank = rank;
          best = i;
        }
      }
      if (best < 0) return symbols;
      symbols.splice(best, 2, symbols[best]! + symbols[best + 1]!);
    }
  }
}

function pairKey(a: string, b: string): string {
  return `${a}\u0000${b}`;
}

function splitMerge(merge: string): [string, string] {
  const space = merge.indexOf(' ');
  if (space < 0) throw new Error(`malformed BPE merge: ${JSON.stringify(merge)}`);
  return [merge.slice(0, space), merge.slice(space + 1)];
}

export function loadBpe(json: BpeFile): BpeTokenizer {
  return new BpeTokenizer(json);
}