Spaces:
Sleeping
Sleeping
File size: 4,218 Bytes
3afc977 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """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"))
|