tokenlib.py

"""Run gpt-tokenizer's browser bundle under QuickJS so the build step and the
browser agree exactly on token boundaries. QuickJS has no TextDecoder/TextEncoder,
so we polyfill both before loading the bundle."""

import json
import os
import sys

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

import quickjs  # noqa: E402

POLYFILL = r"""
var globalThis = this;
function TextDecoder(enc){}
TextDecoder.prototype.decode = function(bytes){
  var out = "", i = 0, n = bytes.length;
  while (i < n) {
    var c = bytes[i++];
    if (c < 0x80) out += String.fromCharCode(c);
    else if (c < 0xE0) out += String.fromCharCode(((c & 0x1F) << 6) | (bytes[i++] & 0x3F));
    else if (c < 0xF0) out += String.fromCharCode(((c & 0x0F) << 12) | ((bytes[i++] & 0x3F) << 6) | (bytes[i++] & 0x3F));
    else {
      var cp = ((c & 0x07) << 18) | ((bytes[i++] & 0x3F) << 12) | ((bytes[i++] & 0x3F) << 6) | (bytes[i++] & 0x3F);
      cp -= 0x10000;
      out += String.fromCharCode(0xD800 + (cp >> 10), 0xDC00 + (cp & 0x3FF));
    }
  }
  return out;
};
function TextEncoder(){}
TextEncoder.prototype.encode = function(str){
  var out = [], i = 0;
  while (i < str.length) {
    var cp = str.codePointAt(i);
    i += cp > 0xFFFF ? 2 : 1;
    if (cp < 0x80) out.push(cp);
    else if (cp < 0x800) out.push(0xC0 | (cp >> 6), 0x80 | (cp & 0x3F));
    else if (cp < 0x10000) out.push(0xE0 | (cp >> 12), 0x80 | ((cp >> 6) & 0x3F), 0x80 | (cp & 0x3F));
    else out.push(0xF0 | (cp >> 18), 0x80 | ((cp >> 12) & 0x3F), 0x80 | ((cp >> 6) & 0x3F), 0x80 | (cp & 0x3F));
  }
  return new Uint8Array(out);
};
"""


class Tokenizer:
    def __init__(self, encoding="o200k_base"):
        self.encoding = encoding
        self.ctx = quickjs.Context()
        self.ctx.set_memory_limit(1 << 30)
        self.ctx.set_max_stack_size(1 << 22)
        self.ctx.eval(POLYFILL)
        path = os.path.join(VENDOR, "gpt-tokenizer", f"{encoding}.js")
        with open(path, encoding="utf-8") as fh:
            self.ctx.eval(fh.read())
        self.ctx.eval(f"var T = globalThis.GPTTokenizer_{encoding};")
        self.ctx.eval("""
        function tokenize(s){
          var ids = T.encode(s);
          var out = [];
          for (var i = 0; i < ids.length; i++) out.push([ids[i], T.decode([ids[i]])]);
          return JSON.stringify(out);
        }
        """)

    def tokens(self, text):
        """-> list of {id, text} in order."""
        self.ctx.set("_input", text)
        raw = self.ctx.eval("tokenize(_input)")
        return [{"id": i, "text": t} for i, t in json.loads(raw)]

    def count(self, text):
        self.ctx.set("_input", text)
        return int(self.ctx.eval("T.encode(_input).length"))

    @property
    def vocab_size(self):
        return int(self.ctx.eval("T.vocabularySize"))