| |
| import numpy as np |
| try: |
| from numba import njit |
| except ImportError: |
|
|
| def njit(*a, **k): |
| return (lambda f: f) if not a or not callable(a[0]) else a[0] |
| C_OTHER, C_ALNUM, C_SPACE, C_TAB, C_NL, C_DQ, C_SQ, C_HASH, C_BSL, C_OPEN, C_CLOSE = range(11) |
| T_NORMAL, T_RESET, T_FIM_SUF, T_FIM_MID = (0, 1, 2, 3) |
|
|
| class Tables: |
|
|
| def __init__(self, tok): |
| V = tok.get_vocab_size() |
| specials = {t.content for t in tok.get_added_tokens_decoder().values()} if hasattr(tok, 'get_added_tokens_decoder') else set() |
| bs = list(range(ord('!'), ord('~') + 1)) + list(range(ord('¡'), ord('¬') + 1)) + list(range(ord('®'), ord('ÿ') + 1)) |
| cs, n = (bs[:], 0) |
| for b in range(256): |
| if b not in bs: |
| bs.append(b) |
| cs.append(256 + n) |
| n += 1 |
| dec = {chr(c): b for b, c in zip(bs, cs)} |
| off, ln, flag, ev = (np.zeros(V, np.int64), np.zeros(V, np.int32), np.zeros(V, np.int8), []) |
| for i in range(V): |
| s = tok.id_to_token(i) or '' |
| if s in specials or (s.startswith('<|') and s.endswith('|>')): |
| flag[i] = {'<|fim_suffix|>': T_FIM_SUF, '<|fim_middle|>': T_FIM_MID}.get(s, T_RESET) |
| off[i], ln[i] = (len(ev), 0) |
| continue |
| raw = bytes((dec.get(ch, 63) for ch in s)) |
| off[i] = len(ev) |
| for b in raw: |
| c = chr(b) |
| ev.append(C_ALNUM if c.isalnum() or c == '_' or b >= 128 else C_SPACE if c == ' ' else C_TAB if c == '\t' else C_NL if c == '\n' else C_DQ if c == '"' else C_SQ if c == "'" else C_HASH if c == '#' else C_BSL if c == '\\' else C_OPEN if c in '([{' else C_CLOSE if c in ')]}' else C_OTHER) |
| ln[i] = len(raw) |
| pref = np.zeros(max(1, len(ev)), np.int8) |
| for i in range(V): |
| if ln[i]: |
| raw = bytes((dec.get(ch, 63) for ch in tok.id_to_token(i) or '')) |
| for j, b in enumerate(raw): |
| pref[off[i] + j] = 1 if chr(b) in 'frbuFRBU' else 0 |
| self.arrays = (off, ln, flag, np.asarray(ev, np.int8), pref) |
|
|
| @njit(cache=True) |
| def features(ids, off, ln, flag, ev, pref): |
| n = ids.shape[0] |
| out = np.zeros(n, np.int16) |
| mode = 0 |
| q = 0 |
| qrun = 0 |
| opened = 0 |
| esc = 0 |
| depth = 0 |
| indent = 0 |
| ls = 1 |
| run = 0 |
| run_pref = 1 |
| s_mode = s_depth = s_indent = s_ls = 0 |
| for t in range(n): |
| i = ids[t] |
| f = flag[i] |
| if f == 1: |
| mode = q = qrun = opened = esc = depth = indent = run = 0 |
| ls = 1 |
| run_pref = 1 |
| elif f == 2: |
| s_mode, s_depth, s_indent, s_ls = (mode, depth, indent, ls) |
| mode = q = qrun = opened = esc = depth = indent = run = 0 |
| ls = 1 |
| run_pref = 1 |
| elif f == 3: |
| mode, depth, indent, ls = (s_mode, s_depth, s_indent, s_ls) |
| q = qrun = opened = esc = run = 0 |
| run_pref = 1 |
| else: |
| o = off[i] |
| for j in range(ln[i]): |
| c = ev[o + j] |
| if c == 4: |
| if mode == 1 or mode == 3: |
| mode = 0 |
| opened = qrun = esc = run = 0 |
| run_pref = 1 |
| ls = 1 |
| indent = 0 |
| continue |
| if ls == 1: |
| if c == 2: |
| indent += 1 |
| continue |
| if c == 3: |
| indent += 4 |
| continue |
| ls = 0 |
| if mode == 0: |
| if c == 1: |
| run += 1 |
| if pref[o + j] == 0: |
| run_pref = 0 |
| else: |
| if c == 5 or c == 6: |
| if c == 5 or run == 0 or (run_pref == 1 and run <= 2): |
| mode, q, opened, qrun = (1, c, 1, 1) |
| elif c == 7: |
| mode = 3 |
| elif c == 9: |
| depth += 1 |
| elif c == 10: |
| depth = depth - 1 if depth > 0 else 0 |
| run = 0 |
| run_pref = 1 |
| elif mode == 1: |
| if esc == 1: |
| esc = 0 |
| opened = 0 |
| elif c == q: |
| if opened == 1: |
| qrun += 1 |
| if qrun == 3: |
| mode, opened, qrun = (2, 0, 0) |
| else: |
| mode, qrun = (0, 0) |
| elif opened == 1 and qrun == 2: |
| mode, opened, qrun = (0, 0, 0) |
| if c == 9: |
| depth += 1 |
| elif c == 10: |
| depth = depth - 1 if depth > 0 else 0 |
| elif c == 7: |
| mode = 3 |
| else: |
| opened = 0 |
| qrun = 0 |
| if c == 8: |
| esc = 1 |
| elif mode == 2: |
| if esc == 1: |
| esc = 0 |
| elif c == 8: |
| esc = 1 |
| qrun = 0 |
| elif c == q: |
| qrun += 1 |
| if qrun == 3: |
| mode, qrun = (0, 0) |
| else: |
| qrun = 0 |
| lvl = indent // 4 |
| if lvl > 15: |
| lvl = 15 |
| d = depth if depth < 15 else 15 |
| m = 1 if mode == 1 else 2 if mode == 2 else 3 if mode == 3 else 0 |
| out[t] = lvl | d << 4 | m << 8 | ls << 10 |
| return out |
|
|
| def describe(v): |
| v = int(v) |
| return dict(indent=v & 15, depth=v >> 4 & 15, mode=('code', 'chaîne', 'triple', 'comm.')[v >> 8 & 3], ls=v >> 10 & 1) |
|
|
