makebundle.py

"""Build a browser bundle for an encoding from gpt-tokenizer's CommonJS source.

Exists because upstream's dist/cl100k_base.js is mislabeled — it emits o200k
tokens. Rather than drop the encoding, we bundle the correct source ourselves.

The module list is not guessed: we load the encoding through the require() shim
and take exactly the modules it actually pulled in. The result is verified
against the reference before it is written.
"""

import json
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "pylib"))
ROOT = os.path.abspath(os.path.join(HERE, ".."))

import quickjs  # noqa: E402

from cjsload import CJS, Reference  # noqa: E402
from tokenlib import POLYFILL  # noqa: E402

OUT_DIR = os.path.join(ROOT, "vendor", "gpt-tokenizer")

LOADER_HEAD = """\
/* gpt-tokenizer %(enc)s — bundled from the package's CommonJS source (MIT).
 *
 * Built by sweedworks.com's own build step, not by upstream. Upstream's
 * prebuilt dist/%(enc)s.js in v3.4.0 is mislabeled: it declares this global
 * but emits o200k_base token IDs. This bundle is checked against the reference
 * implementation over a corpus before shipping. See /.build/makebundle.py.
 *
 * Source: https://github.com/niieani/gpt-tokenizer  (MIT)
 */
(function (global) {
  "use strict";
  var __mods = {};
  var __cache = {};

  function normalize(p) {
    var parts = p.split("/"), out = [];
    for (var i = 0; i < parts.length; i++) {
      if (parts[i] === "." || parts[i] === "") continue;
      if (parts[i] === "..") out.pop(); else out.push(parts[i]);
    }
    return out.join("/");
  }

  function resolve(base, req) {
    var p = req.charAt(0) === "." ? normalize(base + "/" + req) : normalize(req);
    var cands = [p, p + ".js", p + "/index.js"];
    for (var i = 0; i < cands.length; i++) {
      if (Object.prototype.hasOwnProperty.call(__mods, cands[i])) return cands[i];
    }
    throw new Error("gpt-tokenizer bundle: cannot resolve " + req + " from " + base);
  }

  function req_(base, request) {
    var path = resolve(base, request);
    if (__cache[path]) return __cache[path].exports;
    var mod = { exports: {} };
    __cache[path] = mod;
    var dir = path.indexOf("/") < 0 ? "" : path.substring(0, path.lastIndexOf("/"));
    __mods[path](mod.exports, function (r) { return req_(dir, r); }, mod);
    return mod.exports;
  }

  function def(path, fn) { __mods[path] = fn; }

"""

LOADER_TAIL = """
  global.GPTTokenizer_%(enc)s = req_("", "./encoding/%(enc)s.js");
})(typeof globalThis !== "undefined" ? globalThis :
   typeof self !== "undefined" ? self : this);
"""


def module_list(encoding):
    """Exactly the modules the encoding pulls in, in load order."""
    ref = Reference(encoding)
    paths = json.loads(ref.ctx.eval("JSON.stringify(Object.keys(__cache))"))
    return [(os.path.relpath(p, CJS).replace(os.sep, "/"), p) for p in paths]


def build(encoding):
    mods = module_list(encoding)
    parts = [LOADER_HEAD % {"enc": encoding}]
    for key, full in sorted(mods):
        src = open(full, encoding="utf-8").read()
        parts.append(f'  def({json.dumps(key)}, function (exports, require, module) {{\n')
        parts.append(src)
        parts.append("\n  });\n\n")
    parts.append(LOADER_TAIL % {"enc": encoding})
    return "".join(parts), len(mods)


CORPUS = [
    "strawberry", " strawberry", "hello world", "", " ", "\n\n", "\t",
    "1234567890", "1,234,567,890", "127.0.0.1",
    "def total(items):\n    return sum(i.price for i in items)\n",
    "こんにちは世界", "人人生而自由", "안녕하세요", "Привет, мир",
    "مرحبا بالعالم", "नमस्ते दुनिया", "café", "🍓", "👩‍👩‍👧‍👦",
    "a" * 300, "The tokenizer does not know what a word is.",
]


def verify(encoding, source):
    """The bundle must agree with the reference on every string, and round-trip."""
    ctx = quickjs.Context()
    ctx.set_memory_limit(1 << 30)
    ctx.set_max_stack_size(1 << 22)
    ctx.eval(POLYFILL)
    ctx.eval(source)
    ctx.eval(f"var B = globalThis.GPTTokenizer_{encoding};")
    if ctx.eval("typeof B.encode") != "function":
        return ["bundle exposes no encode()"]

    ref = Reference(encoding)
    problems = []
    for text in CORPUS:
        ctx.set("_s", text)
        got = json.loads(ctx.eval("JSON.stringify(B.encode(_s))"))
        want = ref.ids(text)
        if got != want:
            problems.append(f"{text!r}: bundle {got[:8]} != reference {want[:8]}")
            continue
        if ctx.eval("B.decode(B.encode(_s))") != text:
            problems.append(f"{text!r}: round trip changed the text")
    return problems


def main():
    encoding = sys.argv[1] if len(sys.argv) > 1 else "cl100k_base"
    source, n = build(encoding)

    problems = verify(encoding, source)
    if problems:
        print(f"FAIL  {encoding}: bundle does not match reference")
        for p in problems[:10]:
            print(f"  {p}")
        return 1

    out = os.path.join(OUT_DIR, f"{encoding}.js")
    with open(out, "w", encoding="utf-8") as fh:
        fh.write(source)

    import gzip
    gz = len(gzip.compress(source.encode()))
    print(f"PASS  {encoding}: {n} modules, matches reference on {len(CORPUS)} strings")
    print(f"      wrote vendor/gpt-tokenizer/{encoding}.js "
          f"({len(source) / 1e6:.2f} MB raw, {gz / 1024:.0f} KB gzipped)")
    return 0


if __name__ == "__main__":
    sys.exit(main())