ml/tokenizer.py

"""Tokenizer over pure-Lua source. Two modes, one interface.

- char  : one token per character (small vocab, but the model must babble code
          character by character — weak at scale).
- lua   : one token per Lua lexeme; whitespace is DROPPED and decode re-joins
          lexemes with single spaces. Lua is whitespace-insensitive, so the
          space-joined program executes identically to the original — verification
          by execution still holds. This shortens sequences, shrinks the vocab,
          removes indentation babble, and (spaces between every token) eliminates
          the `--` comment hazard for free.

Same interface either way: encode -> list[int], decode -> str, special-token ids,
itos/stoi, vocab_size. Mode is stored in the saved JSON so old checkpoints load
unchanged (absent mode defaults to char).
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field

SPECIALS = [
    "<pad>",  # 0: padding outside content; ignored by loss and attention
    "<bos>",
    "<eos>",
    "<mask>",  # diffusion: a position to be denoised
    "<gap>",   # (legacy char-canvas padding; unused by block diffusion)
    "<pre>",   # AR-FIM: prefix marker
    "<suf>",   # AR-FIM: suffix marker
    "<mid>",   # AR-FIM: middle (hole) marker
    "<unk>",   # token-level: lexeme unseen at build time
]

# Lua lexer: meaningful lexemes only (whitespace is dropped; see module docstring).
_LUA_RE = re.compile(
    r"""
      (?P<ws>[ \t\r\n]+)
    | (?P<name>[A-Za-z_][A-Za-z0-9_]*)
    | (?P<num>[0-9]+\.?[0-9]*)
    | (?P<str>"(?:[^"\\]|\\.)*")
    | (?P<op>==|~=|<=|>=|//|\.\.|::|[-+*/%<>=(){}\[\],;.\#:^&|])
    | (?P<other>\S)
    """,
    re.VERBOSE | re.DOTALL,
)


def lua_lex(text: str) -> list[str]:
    """Lexemes only, no whitespace tokens."""
    return [m.group(0) for m in _LUA_RE.finditer(text) if m.lastgroup != "ws"]


def _is_special(s: str) -> bool:
    return len(s) > 1 and s.startswith("<") and s.endswith(">")


@dataclass
class Tokenizer:
    stoi: dict
    itos: dict
    mode: str = "char"

    @property
    def vocab_size(self) -> int:
        return len(self.stoi)

    @property
    def pad_id(self) -> int: return self.stoi["<pad>"]
    @property
    def bos_id(self) -> int: return self.stoi["<bos>"]
    @property
    def eos_id(self) -> int: return self.stoi["<eos>"]
    @property
    def mask_id(self) -> int: return self.stoi["<mask>"]
    @property
    def gap_id(self) -> int: return self.stoi["<gap>"]
    @property
    def pre_id(self) -> int: return self.stoi["<pre>"]
    @property
    def suf_id(self) -> int: return self.stoi["<suf>"]
    @property
    def mid_id(self) -> int: return self.stoi["<mid>"]

    def _split(self, text: str) -> list[str]:
        return lua_lex(text) if self.mode == "lua" else list(text)

    def encode(self, text: str) -> list[int]:
        unk = self.stoi.get("<unk>")
        out = []
        for u in self._split(text):
            i = self.stoi.get(u, unk)
            if i is not None:
                out.append(i)
        return out

    def decode(self, ids) -> str:
        sep = " " if self.mode == "lua" else ""
        toks = [self.itos[int(i)] for i in ids if not _is_special(self.itos[int(i)])]
        return sep.join(toks)

    @classmethod
    def build(cls, texts: list[str], mode: str = "char") -> "Tokenizer":
        split = lua_lex if mode == "lua" else (lambda t: list(t))
        units = set()
        for t in texts:
            units.update(split(t))
        ordered = SPECIALS + sorted(units)
        stoi = {tok: i for i, tok in enumerate(ordered)}
        itos = {i: tok for tok, i in stoi.items()}
        return cls(stoi=stoi, itos=itos, mode=mode)

    def save(self, path: str) -> None:
        with open(path, "w") as f:
            json.dump({"mode": self.mode,
                       "itos": {str(k): v for k, v in self.itos.items()}}, f)

    @classmethod
    def load(cls, path: str) -> "Tokenizer":
        with open(path) as f:
            d = json.load(f)
        itos = {int(k): v for k, v in d["itos"].items()}
        stoi = {v: k for k, v in itos.items()}
        return cls(stoi=stoi, itos=itos, mode=d.get("mode", "char"))