Upload scripts/prep_base.py with huggingface_hub

scripts/prep_base.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+"""Mistral-Small-3.1-24B-Base-2503 -> thrasher base: strip vision, claim ChatML.
+
+    python prep_base.py --src <base snapshot dir> --out <dir> [--init copy|mean|none]
+    python prep_base.py --src <dir with tokenizer jsons> --out <dir> --tokenizer-only
+
+One streaming pass over the shards (CPU-only, ~one tensor in memory at a time;
+fine on the box or locally). What it does, and why:
+
+  WEIGHTS  drop `vision_tower.*` + `multi_modal_projector.*` (222 keys, ~0.9 GB),
+           strip the `language_model.` prefix off the remaining 363 keys so the
+           result is a plain published-shape MistralForCausalLM dict
+           (lm_head.weight, model.embed_tokens.weight, model.layers.*).
+           lm_head is genuinely untied in the source — verified against
+           model.safetensors.index.json 2026-08-26.
+
+  TOKENS   rename <SPECIAL_20> -> <|im_start|> and <SPECIAL_21> -> <|im_end|>
+           in tokenizer.json (added_tokens + model.vocab), tokenizer_config.json
+           and special_tokens_map.json (recursive string replace — the slot
+           names also appear in additional_special_tokens). eos_token becomes
+           <|im_end|>; bos stays <s> (the tokenizer.json post_processor
+           auto-prepends it on encode, so the jinja emits no BOS and training/
+           GGUF/MLX all see the same bytes).
+
+  INIT     the claimed rows never appeared in pretraining data, so their
+           embed/lm_head rows are dead (the Qwen3-Base never-stops lesson: a
+           dead <eos> lm_head row means the model cannot terminate, and our
+           SFT is ~1B tokens, not Hermes's 60B — don't gamble on cold-claiming).
+           --init copy (default): <|im_start|> rows <- <s> rows, <|im_end|>
+           rows <- </s> rows, embed AND head. Same-model semantic donors, the
+           graft half of the qwen3 graft+heal recipe; SFT is the heal.
+           --init mean: mean of all live rows. --init none: report only.
+           A liveness report (row norms vs live-row stats) prints either way.
+
+  CONFIG   text_config promoted to top level, architectures=[MistralForCausalLM],
+           model_type=mistral, tie_word_embeddings=false explicit. generation_
+           config eos_token_id=[21]. preprocessor/processor configs not copied.
+
+  TEMPLATE thrasher.jinja written to chat_template.jinja AND embedded in
+           tokenizer_config.json (consolidate_ckpt pattern — MLX/GGUF read it
+           from there and the base ships none).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+import sys
+from pathlib import Path
+
+CLAIMS = {  # id -> (old, new, donor_id)
+    20: ("<SPECIAL_20>", "<|im_start|>", 1),   # donor <s>
+    21: ("<SPECIAL_21>", "<|im_end|>", 2),     # donor </s>
+}
+
+# The published tokenizer.json carries a BROKEN pre-tokenizer regex (the
+# known Mistral conversion bug — transformers warns and offers
+# fix_mistral_regex=True at load time, but that fix lives in-memory only).
+# Training and GGUF/MLX read tokenizer.json directly, so the real Tekken
+# pattern must be baked into the file. FIXED is byte-identical to what
+# transformers 5.x fix_mistral_regex produces (diffed on-box 2026-08-26):
+# case-aware word pattern, single-digit \p{N} splits, [\r\n/]* punct tail.
+BROKEN_REGEX = (r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+"
+                r"|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+")
+FIXED_REGEX = (r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*"
+               r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]+|[^\r\n\p{L}\p{N}]?"
+               r"[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*"
+               r"|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+")
+DROP_PREFIXES = ("vision_tower.", "multi_modal_projector.")
+LM_PREFIX = "language_model."
+ROW_KEYS = ("model.embed_tokens.weight", "lm_head.weight")
+
+
+def replace_deep(obj, mapping: dict[str, str]):
+    if isinstance(obj, str):
+        return mapping.get(obj, obj)
+    if isinstance(obj, list):
+        return [replace_deep(x, mapping) for x in obj]
+    if isinstance(obj, dict):
+        return {k: replace_deep(v, mapping) for k, v in obj.items()}
+    return obj
+
+
+def jload(p: Path):
+    with open(p) as f:
+        return json.load(f)
+
+
+def jdump(obj, p: Path):
+    with open(p, "w") as f:
+        json.dump(obj, f, indent=2, ensure_ascii=False)
+        f.write("\n")
+
+
+def prep_tokenizer(src: Path, out: Path, template_text: str) -> None:
+    strmap = {old: new for old, new, _ in CLAIMS.values()}
+
+    tj = jload(src / "tokenizer.json")
+    renamed = 0
+    for tok in tj["added_tokens"]:
+        if tok["content"] in strmap:
+            tok["content"] = strmap[tok["content"]]
+            renamed += 1
+    vocab = tj["model"]["vocab"]
+    for old, new, _ in CLAIMS.values():
+        assert old in vocab, f"{old} not in vocab — wrong base?"
+        assert new not in vocab, f"{new} already in vocab"
+        vocab[new] = vocab.pop(old)
+    assert renamed == len(CLAIMS), f"renamed {renamed} added_tokens, expected {len(CLAIMS)}"
+
+    split = tj["pre_tokenizer"]["pretokenizers"][0]["pattern"]
+    assert split["Regex"] == BROKEN_REGEX, "pre_tokenizer not the known-broken pattern — re-diff before baking"
+    split["Regex"] = FIXED_REGEX
+    jdump(tj, out / "tokenizer.json")
+
+    tc = replace_deep(jload(src / "tokenizer_config.json"), strmap)
+    tc["eos_token"] = "<|im_end|>"
+    tc["chat_template"] = template_text
+    jdump(tc, out / "tokenizer_config.json")
+
+    sm = replace_deep(jload(src / "special_tokens_map.json"), strmap)
+    eos = sm.get("eos_token")
+    if isinstance(eos, dict):
+        eos["content"] = "<|im_end|>"
+    else:
+        sm["eos_token"] = "<|im_end|>"
+    jdump(sm, out / "special_tokens_map.json")
+
+    (out / "chat_template.jinja").write_text(template_text)
+
+    # round-trip proof, not guess
+    from tokenizers import Tokenizer
+    tok = Tokenizer.from_file(str(out / "tokenizer.json"))
+    ids = tok.encode("<|im_start|>user\nhi<|im_end|>\n").ids
+    assert ids[0] == 1 and 20 in ids and 21 in ids, f"claim round-trip failed: {ids}"
+    assert tok.decode([20, 21], skip_special_tokens=False) == "<|im_start|><|im_end|>"
+    print(f"[tokenizer] claimed: " + ", ".join(
+        f"{new}={i}" for i, (_, new, _) in CLAIMS.items()))
+    print(f"[tokenizer] round-trip ids for ChatML probe: {ids}")
+
+
+def prep_configs(src: Path, out: Path) -> None:
+    cfg = jload(src / "config.json")
+    text = cfg["text_config"]
+    text.update({
+        "architectures": ["MistralForCausalLM"],
+        "model_type": "mistral",
+        "torch_dtype": cfg.get("torch_dtype", "bfloat16"),
+        "tie_word_embeddings": False,
+        "bos_token_id": 1,
+        "eos_token_id": 21,
+    })
+    jdump(text, out / "config.json")
+
+    gen = {"bos_token_id": 1, "eos_token_id": [21]}
+    if (src / "generation_config.json").exists():
+        g = jload(src / "generation_config.json")
+        g.update(gen)
+        g.pop("pad_token_id", None)
+        gen = g
+    gen["transformers_version"] = None
+    gen = {k: v for k, v in gen.items() if v is not None}
+    jdump(gen, out / "generation_config.json")
+    print("[config] MistralForCausalLM, untied, eos_token_id=[21]")
+
+
+def prep_weights(src: Path, out: Path, init: str) -> None:
+    import torch
+    from safetensors import safe_open
+    from safetensors.torch import save_file
+
+    index = jload(src / "model.safetensors.index.json")
+    wmap = index["weight_map"]
+    shards: dict[str, list[str]] = {}
+    for key, shard in wmap.items():
+        shards.setdefault(shard, []).append(key)
+
+    new_map: dict[str, str] = {}
+    total = 0
+    n_drop = n_keep = 0
+    donor_rows: dict[str, dict[int, torch.Tensor]] = {}  # row_key -> {donor_id: row}
+
+    shard_names = sorted(shards)
+    for si, shard in enumerate(shard_names, 1):
+        out_name = f"model-{si:05d}-of-{len(shard_names):05d}.safetensors"
+        tensors: dict[str, torch.Tensor] = {}
+        with safe_open(src / shard, framework="pt") as f:
+            for key in sorted(shards[shard]):
+                if key.startswith(DROP_PREFIXES):
+                    n_drop += 1
+                    continue
+                assert key.startswith(LM_PREFIX), f"unexpected key {key}"
+                nk = key[len(LM_PREFIX):]
+                t = f.get_tensor(key)
+                if nk in ROW_KEYS:
+                    t = claim_rows(nk, t, init, donor_rows)
+                tensors[nk] = t
+                n_keep += 1
+        if not tensors:
+            continue
+        save_file(tensors, str(out / out_name), metadata={"format": "pt"})
+        for nk, t in tensors.items():
+            new_map[nk] = out_name
+            total += t.numel() * t.element_size()
+        print(f"[weights] {shard} -> {out_name}  ({len(tensors)} tensors)")
+
+    jdump({"metadata": {"total_size": total}, "weight_map": new_map},
+          out / "model.safetensors.index.json")
+    print(f"[weights] kept {n_keep}, dropped {n_drop}, total {total/1e9:.2f} GB")
+    assert n_keep == 363 and n_drop == 222, "key census mismatch vs 2026-08-26 index"
+
+
+def claim_rows(name: str, t, init: str, donor_rows) -> "torch.Tensor":
+    import torch
+    live = t[1000:]  # rows past the control block are all trained BPE tokens
+    live_norm = live.float().norm(dim=1)
+    print(f"[liveness] {name}: live rows norm mean {live_norm.mean():.4f} "
+          f"(p5 {live_norm.quantile(0.05):.4f})")
+    for tid, (_, new, donor) in CLAIMS.items():
+        print(f"[liveness]   row {tid} ({new}): norm {t[tid].float().norm():.4f}, "
+              f"donor row {donor}: {t[donor].float().norm():.4f}")
+    if init == "none":
+        return t
+    t = t.clone()
+    for tid, (_, _, donor) in CLAIMS.items():
+        t[tid] = t[donor] if init == "copy" else live.float().mean(0).to(t.dtype)
+    return t
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--src", required=True, type=Path)
+    ap.add_argument("--out", required=True, type=Path)
+    ap.add_argument("--template", type=Path,
+                    default=Path(__file__).parent.parent / "configs" / "thrasher.jinja")
+    ap.add_argument("--init", choices=("copy", "mean", "none"), default="copy")
+    ap.add_argument("--tokenizer-only", action="store_true")
+    args = ap.parse_args()
+
+    args.out.mkdir(parents=True, exist_ok=True)
+    template_text = args.template.read_text()
+
+    prep_tokenizer(args.src, args.out, template_text)
+    if args.tokenizer_only:
+        print("[done] tokenizer-only")
+        return
+    prep_configs(args.src, args.out)
+    prep_weights(args.src, args.out, args.init)
+    print("[done] prepped base at", args.out)
+
+
+if __name__ == "__main__":
+    main()