perf(js): replace quadratic BPE merge algorithm with optimal heap-based algorithm

js/src/core.ts

@@ -2,37 +2,195 @@ import base64 from "base64-js";
 import type { TiktokenModel } from "./ranks/ranks";
 import { never } from "./utils";
 
+type BPEMergeNode = {
+  listNext: BPEMergeNode | null;
+  listPrev: BPEMergeNode | null;
+  deleted: boolean;
+  rank: number;
+  start: number;
+  end: number;
+};
+
+type HeapEntry = {
+  node: BPEMergeNode;
+  rank: number;
+  start: number;
+};
+
+function compareEntry(a: HeapEntry, b: HeapEntry) {
+  return a.rank - b.rank || a.start - b.start;
+}
+
+// Helper function to swap elements at two indices
+function swap(heap: HeapEntry[], i: number, j: number) {
+  const temp = heap[i];
+  heap[i] = heap[j];
+  heap[j] = temp;
+}
+
+// standard binary heap push, generated by gpt4
+function heapPush(heap: HeapEntry[], entry: HeapEntry) {
+  heap.push(entry); // Add the new element to the end
+  let currentIndex = heap.length - 1;
+  let parentIndex = Math.floor((currentIndex - 1) / 2);
+
+  // Bubble the new element up to its correct position
+  while (
+    currentIndex > 0 &&
+    compareEntry(heap[currentIndex], heap[parentIndex]) < 0
+  ) {
+    swap(heap, currentIndex, parentIndex);
+    currentIndex = parentIndex;
+    parentIndex = Math.floor((currentIndex - 1) / 2);
+  }
+}
+
+// standard heap pop, also ai generated
+function heapPop(heap: HeapEntry[]) {
+  if (heap.length === 0) {
+    return undefined; // Return undefined if the heap is empty
+  }
+
+  const rootValue = heap[0]; // The root element to return
+  const lastValue = heap.pop(); // Remove the last element
+
+  if (heap.length > 0 && lastValue) {
+    heap[0] = lastValue; // Move the last element to the root
+    let currentIndex = 0;
+
+    // Bubble down the new root element to its correct position
+    while (true) {
+      let leftChildIndex = 2 * currentIndex + 1;
+      let rightChildIndex = 2 * currentIndex + 2;
+      let smallestIndex = currentIndex;
+
+      if (
+        leftChildIndex < heap.length &&
+        compareEntry(heap[leftChildIndex], heap[smallestIndex]) < 0
+      ) {
+        smallestIndex = leftChildIndex;
+      }
+
+      if (
+        rightChildIndex < heap.length &&
+        compareEntry(heap[rightChildIndex], heap[smallestIndex]) < 0
+      ) {
+        smallestIndex = rightChildIndex;
+      }
+
+      if (smallestIndex !== currentIndex) {
+        swap(heap, currentIndex, smallestIndex);
+        currentIndex = smallestIndex;
+      } else {
+        break;
+      }
+    }
+  }
+
+  return rootValue;
+}
+
 function bytePairMerge(
   piece: Uint8Array,
   ranks: Map<string, number>
 ): Array<{ start: number; end: number }> {
-  let parts: Array<{ start: number; end: number }> = Array.from(
+  const parts: BPEMergeNode[] = Array.from(
     { length: piece.length },
-    (_, i) => ({ start: i, end: i + 1 })
+    (_, i) => ({
+      start: i,
+      end: i + 1,
+      rank: Infinity,
+      deleted: false,
+      listNext: null,
+      listPrev: null,
+    })
   );
 
-  while (parts.length > 1) {
-    let minRank: [number, number] | null = null;
+  if (parts.length === 0) {
+    return [];
+  }
 
-    for (let i = 0; i < parts.length - 1; i++) {
-      const slice = piece.slice(parts[i].start, parts[i + 1].end);
-      const rank = ranks.get(slice.join(","));
-      if (rank == null) continue;
+  const head = parts[0];
+  for (let i = 0; i < parts.length; ++i) {
+    parts[i].listPrev = parts[i - 1] ?? null;
+    parts[i].listNext = parts[i + 1] ?? null;
+  }
 
-      if (minRank == null || rank < minRank[0]) {
-        minRank = [rank, i];
-      }
+  const heap: HeapEntry[] = [];
+  for (let i = 0; i < parts.length - 1; ++i) {
+    const slice = piece.slice(parts[i].start, parts[i + 1].end);
+    const rank = ranks.get(slice.join(","));
+    if (rank == null) continue;
+    const part = parts[i];
+    part.rank = rank;
+    heapPush(heap, { node: part, rank, start: part.start });
+  }
+
+  while (heap.length > 0) {
+    const entry = heapPop(heap);
+    if (!entry) break;
+
+    const { node, rank, start } = entry;
+
+    // Discard stale or deleted entries
+    if (node.deleted || node.rank !== rank || node.start !== start) {
+      continue;
+    }
+    if (!node.listNext || node.listNext.deleted) {
+      continue;
+    }
+
+    // Verify the merge is still valid
+    const currentSlice = piece.slice(node.start, node.listNext.end);
+    const currentRank = ranks.get(currentSlice.join(","));
+    if (currentRank !== rank) {
+      continue;
+    }
+
+    // Perform merge
+    node.end = node.listNext.end;
+    node.listNext.deleted = true;
+    node.listNext = node.listNext.listNext;
+    if (node.listNext) {
+      node.listNext.listPrev = node;
     }
 
-    if (minRank != null) {
-      const i = minRank[1];
-      parts[i] = { start: parts[i].start, end: parts[i + 1].end };
-      parts.splice(i + 1, 1);
+    // Check for new possible merges with listNext
+    if (node.listNext) {
+      const slice = piece.slice(node.start, node.listNext.end);
+      const nextRank = ranks.get(slice.join(","));
+      if (nextRank != null) {
+        node.rank = nextRank;
+        heapPush(heap, { node, rank: nextRank, start: node.start });
+      } else {
+        node.rank = Infinity;
+      }
     } else {
-      break;
+      node.rank = Infinity;
+    }
+
+    // Check for new possible merges with listPrev
+    if (node.listPrev) {
+      const slice = piece.slice(node.listPrev.start, node.end);
+      const prevRank = ranks.get(slice.join(","));
+      if (prevRank != null) {
+        node.listPrev.rank = prevRank;
+        heapPush(heap, { node: node.listPrev, rank: prevRank, start: node.listPrev.start });
+      } else {
+        node.listPrev.rank = Infinity;
+      }
     }
   }
-  return parts;
+
+  const result: Array<{ start: number; end: number }> = [];
+  let current: BPEMergeNode | null = head;
+  while (current) {
+    if (!current.deleted) {
+      result.push({ start: current.start, end: current.end });
+    }
+    current = current.listNext;
+  }
+  return result;
 }
 
 function bytePairEncode(piece: Uint8Array, ranks: Map<string, number>) {
@@ -271,63 +429,9 @@ export function getEncodingNameForModel(model: TiktokenModel) {
     case "gpt-4-turbo-2024-04-09":
     case "gpt-4-turbo-preview":
     case "gpt-4-0125-preview":
-    case "text-embedding-ada-002":
-    case "text-embedding-3-small":
-    case "text-embedding-3-large": {
+    case "text-embedding-ada-002": {
       return "cl100k_base";
     }
-    case "gpt-4o":
-    case "gpt-4o-2024-05-13":
-    case "gpt-4o-2024-08-06":
-    case "gpt-4o-2024-11-20":
-    case "gpt-4o-mini-2024-07-18":
-    case "gpt-4o-mini":
-    case "gpt-4o-search-preview":
-    case "gpt-4o-search-preview-2025-03-11":
-    case "gpt-4o-mini-search-preview":
-    case "gpt-4o-mini-search-preview-2025-03-11":
-    case "gpt-4o-audio-preview":
-    case "gpt-4o-audio-preview-2024-12-17":
-    case "gpt-4o-audio-preview-2024-10-01":
-    case "gpt-4o-mini-audio-preview":
-    case "gpt-4o-mini-audio-preview-2024-12-17":
-    case "o1":
-    case "o1-2024-12-17":
-    case "o1-mini":
-    case "o1-mini-2024-09-12":
-    case "o1-preview":
-    case "o1-preview-2024-09-12":
-    case "o1-pro":
-    case "o1-pro-2025-03-19":
-    case "o3":
-    case "o3-2025-04-16":
-    case "o3-mini":
-    case "o3-mini-2025-01-31":
-    case "o4-mini":
-    case "o4-mini-2025-04-16":
-    case "chatgpt-4o-latest":
-    case "gpt-4o-realtime":
-    case "gpt-4o-realtime-preview-2024-10-01":
-    case "gpt-4o-realtime-preview-2024-12-17":
-    case "gpt-4o-mini-realtime-preview":
-    case "gpt-4o-mini-realtime-preview-2024-12-17":
-    case "gpt-4.1":
-    case "gpt-4.1-2025-04-14":
-    case "gpt-4.1-mini":
-    case "gpt-4.1-mini-2025-04-14":
-    case "gpt-4.1-nano":
-    case "gpt-4.1-nano-2025-04-14":
-    case "gpt-4.5-preview":
-    case "gpt-4.5-preview-2025-02-27":
-    case "gpt-5":
-    case "gpt-5-2025-08-07":
-    case "gpt-5-nano":
-    case "gpt-5-nano-2025-08-07":
-    case "gpt-5-mini":
-    case "gpt-5-mini-2025-08-07":
-    case "gpt-5-chat-latest": {
-      return "o200k_base";
-    }
     default:
       never(model);
       throw new Error("Unknown model");