Confidence-gated ghost text for coding prompts: Kneser-Ney compile, 4.5 MB browser export + zero-dep runtime, held-out precision/keystrokes-saved benchmarks
4113449 verified | """Compile text into a Llama-format autocomplete model. No training. | |
| Every weight is set by a closed-form rule from n-gram counts; nothing is fit by | |
| gradient descent and no other model is used. The result is an exact transformer | |
| implementation of an interpolated n-gram language model (Witten-Bell or | |
| Kneser-Ney) over the source text. | |
| Residual stream layout (hidden_size = 1 + (n_prev + 1) * CODE_DIM + vocab): | |
| [0] constant 1 (acts as a bias, Llama has none) | |
| cur slot (CODE_DIM) code of the current token <- embedding | |
| prev1..n (CODE_DIM) codes of the 1..n previous tokens <- attention | |
| out slot (vocab) next-token logits, scaled by OUT_SCALE <- MLP | |
| Token codes are seeded random unit vectors (nearly orthogonal). Attention head | |
| a (a = 1..n_prev) has constant query/key vectors whose RoPE phases make the | |
| score peak at exactly a positions back; its value copies that token's code into | |
| slot prev_a. Positions before the start fall back to position 0 (BOS). | |
| The MLP has one neuron per kept context (1..n_prev+1 tokens) plus one always-on | |
| unigram neuron. A neuron's gate sums code matches across slots minus (k - 0.5), | |
| so it fires only when all k tokens match. It writes | |
| log p(next | ctx) - log p(next | shorter ctx) into the out slot. All suffixes of | |
| the true context fire together, so the sum telescopes to the log probability | |
| under the longest kept context. | |
| Compile settings live in compile_settings.json ("hf" for this dense model, | |
| "web" for the sparse browser export written by export_web.py). | |
| """ | |
| import json | |
| import math | |
| import zlib | |
| from collections import Counter, defaultdict | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from safetensors.torch import save_file | |
| from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, processors | |
| from transformers import GenerationConfig, LlamaConfig, PreTrainedTokenizerFast | |
| ROOT = Path(__file__).parent | |
| PHRASE_DIR = ROOT / "phrases" # hand-written everyday phrases | |
| DATA_DIR = ROOT / "data" # human-written prompts/replies (scripts/fetch_coding_prompts.py) | |
| SETTINGS_PATH = ROOT / "compile_settings.json" | |
| EOS, UNK = "<|endoftext|>", "<unk>" # EOS doubles as BOS and pad, like GPT-2 | |
| EOS_ID, UNK_ID = 0, 1 | |
| SEED = 1234 | |
| CODE_DIM = 384 | |
| MAX_POS = 128 | |
| ATTN_MARGIN = 30.0 # min attention logit gap (nats) between target and any other position | |
| BETA = 100.0 # gate sharpness; matched neurons sit at z = BETA / 2 | |
| OUT_SCALE = 1e-5 # keeps the out slot tiny so the final RMSNorm scale stays fixed | |
| UNK_LOGIT = -30.0 | |
| ROPE_THETA_CANDIDATES = [2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 1000.0, 10000.0] | |
| NORMALIZER = normalizers.Sequence([normalizers.NFKC(), normalizers.Lowercase()]) | |
| PRE_TOKENIZER = pre_tokenizers.ByteLevel(add_prefix_space=True) | |
| class Settings: | |
| n_prev: int = 3 # previous tokens attended to; context = n_prev + current token | |
| smoothing: str = "wb" # "wb" (Witten-Bell) or "kn" (interpolated Kneser-Ney) | |
| discount: float = 0.75 # Kneser-Ney absolute discount | |
| phrase_weight: float = 5.0 # weight of one hand-written phrase (phrases/*.txt) | |
| prompt_weight: float = 1.0 # weight of one user prompt (data/prompts.txt) | |
| reply_weight: float = 1.0 # weight of one assistant reply sentence (data/replies.txt) | |
| max_vocab: int = 3000 | |
| min_word_count: float = 3 | |
| min_count: float | None = None # None: smallest context count that fits param_target | |
| param_target: int = 400_000_000 | |
| def order(self) -> int: | |
| return self.n_prev + 1 | |
| def load_settings(name: str) -> Settings: | |
| return Settings(**json.loads(SETTINGS_PATH.read_text())[name]) | |
| # ---------- text ---------- | |
| def pieces(text: str) -> list[str]: | |
| return [p for p, _ in PRE_TOKENIZER.pre_tokenize_str(NORMALIZER.normalize_str(text))] | |
| def load_lines(paths: list[Path]) -> list[list[str]]: | |
| seen, out = set(), [] | |
| for path in paths: | |
| for line in path.read_text().splitlines(): | |
| line = line.strip() | |
| if not line or line.startswith("#"): | |
| continue | |
| toks = pieces(line) | |
| if tuple(toks) not in seen: | |
| seen.add(tuple(toks)) | |
| out.append(toks) | |
| return out | |
| def load_phrases() -> list[list[str]]: | |
| return load_lines(sorted(PHRASE_DIR.glob("*.txt"))) | |
| def load_sources() -> dict[str, list[list[str]]]: | |
| return { | |
| "phrases": load_phrases(), | |
| "prompts": load_lines([DATA_DIR / "prompts.txt"]), | |
| "replies": load_lines([DATA_DIR / "replies.txt"]), | |
| } | |
| def heldout_split(lines: list[list[str]], pct: int = 10) -> tuple[list, list]: | |
| """Deterministic split by content hash: (train, test).""" | |
| test = [t for t in lines if zlib.crc32(" ".join(t).encode()) % 100 < pct] | |
| train = [t for t in lines if zlib.crc32(" ".join(t).encode()) % 100 >= pct] | |
| return train, test | |
| # ---------- counting ---------- | |
| def weighted_sources(s: Settings, sources: dict) -> list[tuple[list[list[str]], float]]: | |
| weights = {"phrases": s.phrase_weight, "prompts": s.prompt_weight, "replies": s.reply_weight} | |
| return [(sources[name], w) for name, w in weights.items() if w > 0 and name in sources] | |
| def build_vocab(s: Settings, sources: dict) -> list[str]: | |
| freq = Counter() | |
| for lines, w in weighted_sources(s, sources): | |
| for toks in lines: | |
| for p in toks: | |
| freq[p] += w | |
| by_freq = lambda p: (-freq[p], p) # noqa: E731 | |
| hand = {p for t in sources["phrases"] for p in t} if s.phrase_weight > 0 else set() | |
| rest = sorted((p for p in freq if p not in hand and freq[p] >= s.min_word_count), key=by_freq) | |
| words = hand | set(rest[: max(0, s.max_vocab - 2 - len(hand))]) | |
| return [EOS, UNK] + sorted(words, key=by_freq) | |
| def padded_events(seq: list[int], order: int): | |
| """(contexts of length 0..order, next token) for every position of a sequence.""" | |
| # `order` BOS tokens: the last one is the real BOS position the model sees. | |
| padded = [EOS_ID] * order + seq + [EOS_ID] | |
| for i in range(order - 1, len(padded) - 1): | |
| yield [tuple(padded[i - k + 1 : i + 1]) if k else () for k in range(order + 1)], padded[i + 1] | |
| def count_ngrams(weighted: list[tuple[list[int], float]], order: int) -> list[defaultdict]: | |
| counts = [defaultdict(Counter) for _ in range(order + 1)] | |
| for seq, weight in weighted: | |
| for ctxs, nxt in padded_events(seq, order): | |
| if nxt == UNK_ID: # never learn to predict <unk> | |
| continue | |
| for k, ctx in enumerate(ctxs): | |
| counts[k][ctx][nxt] += weight | |
| return counts | |
| class NGram: | |
| """Interpolated n-gram over the kept contexts; the reference the weights must reproduce. | |
| Witten-Bell: p_k(j) = (c_j + u * p_{k-1}(j)) / (c + u) | |
| Kneser-Ney: p_k(j) = (c_j - d_j) / c + (sum_j d_j / c) * p_{k-1}(j), d_j = min(D, c_j), | |
| with continuation counts below the top order (KenLM-style) | |
| Both have the form p_k(j) = A_j + gamma * p_{k-1}(j), so each context's neuron can write | |
| log p_k(j) - log(gamma * p_{k-1}(j)) for its seen tokens and 0 elsewhere. | |
| """ | |
| def __init__(self, raw: list[dict], keep: list[set], vocab_size: int, s: Settings): | |
| self.s, self.V, self.order = s, vocab_size, s.order | |
| self.support = {ctx: sum(raw[k][ctx].values()) for k in range(1, self.order + 1) for ctx in keep[k]} | |
| if s.smoothing == "kn": | |
| cont = [defaultdict(Counter) for _ in range(self.order)] | |
| for k in range(self.order): | |
| for ctx, table in raw[k + 1].items(): | |
| for j in table: | |
| cont[k][ctx[1:]][j] += 1 | |
| self.tables = [cont[0][()]] + [{c: cont[k][c] for c in keep[k]} for k in range(1, self.order)] | |
| self.tables.append({c: raw[self.order][c] for c in keep[self.order]}) | |
| elif s.smoothing == "wb": | |
| self.tables = [raw[0][()]] + [{c: raw[k][c] for c in keep[k]} for k in range(1, self.order + 1)] | |
| else: | |
| raise ValueError(s.smoothing) | |
| for k in range(2, self.order + 1): | |
| assert all(c[1:] in self.tables[k - 1] for c in self.tables[k]), "backoff chain broken" | |
| uni = np.zeros(vocab_size) | |
| for j, c in self.tables[0].items(): | |
| uni[j] = c | |
| logits = np.log(uni + 1.0) # add-one smoothing | |
| logits[UNK_ID] = UNK_LOGIT | |
| self.base_logits = logits - logits[2:].mean() # constant shift; softmax-invariant | |
| p0 = np.exp(self.base_logits - self.base_logits.max()) | |
| self.p0 = p0 / p0.sum() | |
| self._prob, self._delta, self._stats = {}, {}, {} | |
| def _table_stats(self, ctx: tuple): | |
| if ctx not in self._stats: | |
| table = self.tables[len(ctx)][ctx] | |
| total = sum(table.values()) | |
| if self.s.smoothing == "wb": | |
| u = len(table) | |
| self._stats[ctx] = (table, {j: c / (total + u) for j, c in table.items()}, u / (total + u)) | |
| else: | |
| d = {j: min(self.s.discount, c) for j, c in table.items()} | |
| self._stats[ctx] = (table, {j: (c - d[j]) / total for j, c in table.items()}, sum(d.values()) / total) | |
| return self._stats[ctx] | |
| def prob(self, j: int, ctx: tuple) -> float: | |
| key = (j, ctx) | |
| if key not in self._prob: | |
| if not ctx: | |
| p = self.p0[j] | |
| else: | |
| lower = self.prob(j, ctx[1:]) | |
| if ctx in self.tables[len(ctx)]: | |
| _, a, gamma = self._table_stats(ctx) | |
| p = a.get(j, 0.0) + gamma * lower | |
| else: | |
| p = lower | |
| self._prob[key] = p | |
| return self._prob[key] | |
| def delta(self, ctx: tuple) -> dict[int, float]: | |
| """log p_k - log p_{k-1}, minus the constant log(gamma) shared by unseen tokens.""" | |
| if ctx not in self._delta: | |
| table, a, gamma = self._table_stats(ctx) | |
| out = {} | |
| for j in table: | |
| q = gamma * self.prob(j, ctx[1:]) | |
| out[j] = math.log(a[j] + q) - math.log(q) | |
| self._delta[ctx] = out | |
| return self._delta[ctx] | |
| def delta_arrays(self, ctx: tuple) -> tuple[np.ndarray, np.ndarray]: | |
| d = self.delta(ctx) | |
| return np.fromiter(d.keys(), dtype=np.int64, count=len(d)), np.fromiter(d.values(), dtype=np.float64, count=len(d)) | |
| def contexts(self) -> list[tuple]: | |
| return [ctx for k in range(1, self.order + 1) for ctx in sorted(self.tables[k])] | |
| def window(self, ids: list[int]) -> tuple: | |
| """Last order tokens (oldest .. current); positions before 0 read BOS like the attention heads.""" | |
| t = len(ids) - 1 | |
| return tuple(ids[max(t - a, 0)] for a in range(self.s.n_prev, -1, -1)) | |
| def predict(self, ids: list[int]) -> tuple[np.ndarray, float, int]: | |
| """(logits up to a constant, support of the longest matched context, its length).""" | |
| window = self.window(ids) | |
| logits = self.base_logits.copy() | |
| support, length = 0.0, 0 | |
| for k in range(1, self.order + 1): | |
| ctx = window[-k:] | |
| if ctx in self.tables[k]: | |
| idx, vals = self.delta_arrays(ctx) | |
| logits[idx] += vals | |
| support, length = self.support[ctx], k | |
| return logits, support, length | |
| def reference_logits(self, ids: list[int]) -> torch.Tensor: | |
| return torch.from_numpy(self.predict(ids)[0]) | |
| def compile_tables(s: Settings, sources: dict): | |
| """vocab, NGram and stats for settings s (min_count=None fits the dense param budget).""" | |
| vocab = build_vocab(s, sources) | |
| index = {p: i for i, p in enumerate(vocab)} | |
| weighted = [([index.get(p, UNK_ID) for p in toks], w) for lines, w in weighted_sources(s, sources) for toks in lines] | |
| raw = count_ngrams(weighted, s.order) | |
| forced = set() | |
| if s.phrase_weight > 0: | |
| for toks in sources["phrases"]: | |
| for ctxs, _ in padded_events([index[p] for p in toks], s.order): | |
| forced.update(ctxs[1:]) | |
| totals = {ctx: sum(t.values()) for k in range(1, s.order + 1) for ctx, t in raw[k].items()} | |
| min_count = s.min_count | |
| if min_count is None: | |
| # Smallest count threshold whose kept neurons (forced + total >= threshold) fit the budget. | |
| hist = Counter(0 if ctx in forced else c for ctx, c in totals.items()) | |
| thresholds = sorted(c for c in hist if c > 0) | |
| running, kept_at = hist[0], {} | |
| for c in reversed(thresholds): | |
| running += hist[c] | |
| kept_at[c] = running | |
| fits = [c for c in thresholds if dense_params(len(vocab), s.n_prev, 1 + kept_at[c]) <= s.param_target] | |
| if not fits: | |
| raise SystemExit("hand-written contexts alone exceed param_target; lower max_vocab") | |
| min_count = fits[0] | |
| keep = [set()] + [{c for c in raw[k] if c in forced or totals[c] >= min_count} for k in range(1, s.order + 1)] | |
| ngram = NGram(raw, keep, len(vocab), s) | |
| stats = { | |
| "vocab": len(vocab), | |
| "min_context_count": min_count, | |
| "contexts_seen": len(totals), | |
| "contexts_kept": sum(len(k) for k in keep), | |
| } | |
| return vocab, ngram, stats | |
| # ---------- dense transformer ---------- | |
| def hidden_size(V: int, n_prev: int) -> int: | |
| # transformers validates hidden_size % num_heads == 0; pad with unused zero dims. | |
| raw = 1 + (n_prev + 1) * CODE_DIM + V | |
| return -(-raw // (8 * n_prev)) * (8 * n_prev) | |
| def dense_params(V: int, n_prev: int, neurons: int) -> int: | |
| D, h = hidden_size(V, n_prev), CODE_DIM | |
| return 2 * V * D + (2 * n_prev + 2) * h * D + 3 * neurons * D + 3 * D | |
| def rope_attention_scale(head_dim: int, n_prev: int) -> tuple[float, torch.Tensor, float]: | |
| """Pick rope_theta maximizing the worst-case attention margin; return (theta, inv_freq, q_amp).""" | |
| best = None | |
| for theta in ROPE_THETA_CANDIDATES: | |
| inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.int64).double() / head_dim)) | |
| offsets = torch.arange(-MAX_POS, MAX_POS + 1, dtype=torch.float64) | |
| kernel = torch.cos(offsets[:, None] * inv_freq[None, :]).sum(1) # K(x), index x + MAX_POS | |
| worst = math.inf | |
| for lag in range(1, n_prev + 1): | |
| for t in range(MAX_POS): | |
| s_target = max(t - lag, 0) | |
| xs = torch.tensor([t - s - lag for s in range(t + 1)]) + MAX_POS | |
| scores = kernel[xs] | |
| others = torch.cat([scores[:s_target], scores[s_target + 1 :]]) | |
| if len(others): | |
| worst = min(worst, (scores[s_target] - others.max()).item()) | |
| if best is None or worst > best[0]: | |
| best = (worst, theta, inv_freq) | |
| worst, theta, inv_freq = best | |
| assert worst > 0, "no rope_theta separates positions" | |
| # score = q_amp * sum_i cos(...) / sqrt(head_dim) (key amplitude 1) | |
| q_amp = ATTN_MARGIN * math.sqrt(head_dim) / worst | |
| return theta, inv_freq, q_amp | |
| def build_dense(vocab: list[str], ngram: NGram, s: Settings): | |
| V, m, n_prev = len(vocab), CODE_DIM, s.n_prev | |
| contexts = ngram.contexts() | |
| N = 1 + len(contexts) | |
| CONST, CUR = 0, 1 | |
| prev_slot = [1 + a * m for a in range(n_prev + 1)] # prev_slot[0] is the cur slot | |
| OUT = 1 + (n_prev + 1) * m | |
| n_heads, h = n_prev, m | |
| D = hidden_size(V, n_prev) | |
| gen = torch.Generator().manual_seed(SEED) | |
| codes = torch.randn(V, m, generator=gen, dtype=torch.float64) | |
| codes /= codes.norm(dim=1, keepdim=True) | |
| gram = codes @ codes.T | |
| gram.fill_diagonal_(0) | |
| max_coherence = gram.abs().max().item() | |
| assert max_coherence < 0.4, max_coherence | |
| embed = torch.zeros(V, D, dtype=torch.float64) | |
| embed[:, CONST] = 1.0 | |
| embed[:, CUR : CUR + m] = codes | |
| theta, inv_freq, q_amp = rope_attention_scale(h, n_prev) | |
| half = h // 2 | |
| q_proj = torch.zeros(n_heads * h, D, dtype=torch.float64) | |
| for a in range(n_heads): | |
| lag = a + 1 | |
| q_proj[a * h : a * h + half, CONST] = q_amp * torch.cos(lag * inv_freq) | |
| q_proj[a * h + half : (a + 1) * h, CONST] = -q_amp * torch.sin(lag * inv_freq) | |
| k_proj = torch.zeros(h, D, dtype=torch.float64) | |
| k_proj[:half, CONST] = 1.0 | |
| v_proj = torch.zeros(h, D, dtype=torch.float64) | |
| v_proj[:, CUR : CUR + m] = torch.eye(m, dtype=torch.float64) | |
| o_proj = torch.zeros(D, n_heads * h, dtype=torch.float64) | |
| for a in range(n_heads): | |
| o_proj[prev_slot[a + 1] : prev_slot[a + 1] + m, a * h : (a + 1) * h] = torch.eye(m, dtype=torch.float64) | |
| # The MLP matrices dominate memory, so they are built directly in float32. | |
| gate = torch.zeros(N, D, dtype=torch.float32) | |
| up = torch.zeros(N, D, dtype=torch.float32) | |
| down = torch.zeros(D, N, dtype=torch.float32) | |
| on = (BETA / 2) * torch.sigmoid(torch.tensor(BETA / 2, dtype=torch.float64)).item() # silu(BETA/2) | |
| up[:, CONST] = 1.0 | |
| gate[0, CONST] = BETA / 2 # always-on unigram neuron | |
| down[OUT : OUT + V, 0] = torch.from_numpy(OUT_SCALE * ngram.base_logits / on).float() | |
| for n, ctx in enumerate(contexts, start=1): | |
| k = len(ctx) | |
| for a, tok in enumerate(reversed(ctx)): # a = 0 current, a = 1 previous, ... | |
| gate[n, prev_slot[a] : prev_slot[a] + m] = BETA * codes[tok] | |
| gate[n, CONST] = -BETA * (k - 0.5) | |
| for j, d in ngram.delta(ctx).items(): | |
| down[OUT + j, n] = OUT_SCALE * d / on | |
| lm_head = torch.zeros(V, D, dtype=torch.float64) | |
| lm_head[torch.arange(V), OUT + torch.arange(V)] = 1.0 / OUT_SCALE | |
| # RMSNorm weights undo the norm for the (fixed) input norms: |x|^2 = 2 before | |
| # attention, 1 + 1 + n_prev after it. | |
| attn_norm = math.sqrt(2 / D) | |
| mlp_norm = math.sqrt((2 + n_prev) / D) | |
| state = { | |
| "model.embed_tokens.weight": embed, | |
| "model.layers.0.input_layernorm.weight": torch.full((D,), attn_norm), | |
| "model.layers.0.self_attn.q_proj.weight": q_proj, | |
| "model.layers.0.self_attn.k_proj.weight": k_proj, | |
| "model.layers.0.self_attn.v_proj.weight": v_proj, | |
| "model.layers.0.self_attn.o_proj.weight": o_proj, | |
| "model.layers.0.post_attention_layernorm.weight": torch.full((D,), mlp_norm), | |
| "model.layers.0.mlp.gate_proj.weight": gate, | |
| "model.layers.0.mlp.up_proj.weight": up, | |
| "model.layers.0.mlp.down_proj.weight": down, | |
| "model.norm.weight": torch.full((D,), mlp_norm), | |
| "lm_head.weight": lm_head, | |
| } | |
| state = {k: v.to(torch.float32).contiguous() for k, v in state.items()} | |
| config = LlamaConfig( | |
| vocab_size=V, | |
| hidden_size=D, | |
| intermediate_size=N, | |
| num_hidden_layers=1, | |
| num_attention_heads=n_heads, | |
| num_key_value_heads=1, | |
| head_dim=h, | |
| hidden_act="silu", | |
| max_position_embeddings=MAX_POS, | |
| rms_norm_eps=1e-10, | |
| rope_theta=theta, | |
| tie_word_embeddings=False, | |
| bos_token_id=0, | |
| eos_token_id=0, | |
| pad_token_id=0, | |
| ) | |
| info = { | |
| "hidden": D, | |
| "neurons": N, | |
| "rope_theta": theta, | |
| "max_code_coherence": round(max_coherence, 3), | |
| "params": sum(t.numel() for t in state.values()), | |
| } | |
| return state, config, info | |
| def build_tokenizer(vocab: list[str]) -> PreTrainedTokenizerFast: | |
| tok = Tokenizer(models.WordLevel({p: i for i, p in enumerate(vocab)}, unk_token=UNK)) | |
| tok.normalizer = NORMALIZER | |
| tok.pre_tokenizer = PRE_TOKENIZER | |
| tok.decoder = decoders.ByteLevel() | |
| tok.post_processor = processors.TemplateProcessing(single=f"{EOS} $A", special_tokens=[(EOS, 0)]) | |
| return PreTrainedTokenizerFast( | |
| tokenizer_object=tok, | |
| bos_token=EOS, | |
| eos_token=EOS, | |
| unk_token=UNK, | |
| pad_token=EOS, | |
| clean_up_tokenization_spaces=False, | |
| model_max_length=MAX_POS, | |
| ) | |
| def compile_model(): | |
| """Build the dense ("hf") model in memory (used by the tests for the n-gram reference).""" | |
| s = load_settings("hf") | |
| sources = load_sources() | |
| vocab, ngram, stats = compile_tables(s, sources) | |
| state, config, info = build_dense(vocab, ngram, s) | |
| lines = {name: len(v) for name, v in sources.items()} | |
| return lines, vocab, ngram, state, config, {**stats, **info} | |
| def main() -> None: | |
| lines, vocab, _, state, config, info = compile_model() | |
| assert info["params"] <= 1_000_000_000, info | |
| save_file(state, ROOT / "model.safetensors", metadata={"format": "pt"}) | |
| config.save_pretrained(ROOT) | |
| cfg_path = ROOT / "config.json" | |
| cfg = json.loads(cfg_path.read_text()) | |
| cfg["architectures"] = ["LlamaForCausalLM"] | |
| cfg["rope_theta"] = info["rope_theta"] # transformers 4.x reads this key, 5.x reads rope_parameters | |
| cfg["torch_dtype"] = "float32" | |
| cfg_path.write_text(json.dumps(cfg, indent=2, sort_keys=True) + "\n") | |
| GenerationConfig(bos_token_id=0, eos_token_id=0, pad_token_id=0, do_sample=False, max_new_tokens=16).save_pretrained(ROOT) | |
| build_tokenizer(vocab).save_pretrained(ROOT) | |
| # transformers 5 writes "TokenizersBackend", which 4.x can't resolve; | |
| # PreTrainedTokenizerFast loads on both. 4.x also emits token_type_ids | |
| # by default, which Llama's generate() rejects, and overrides the | |
| # pre-tokenizer's add_prefix_space unless the config states it. | |
| tok_cfg_path = ROOT / "tokenizer_config.json" | |
| tok_cfg = json.loads(tok_cfg_path.read_text()) | |
| tok_cfg["tokenizer_class"] = "PreTrainedTokenizerFast" | |
| tok_cfg["add_prefix_space"] = True | |
| tok_cfg["model_input_names"] = ["input_ids", "attention_mask"] | |
| tok_cfg_path.write_text(json.dumps(tok_cfg, indent=2) + "\n") | |
| print(json.dumps({"lines": lines, **info}, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |