Debut release of Juker-Coder-32B (PyTorch Spectral Gravity-Well Splice + Nullspace-Locked Phase Heal & Full 12-Benchmark Suite)
b3bdb85 verified | #!/usr/bin/env python3 | |
| """ | |
| Merge Juker (31B) Surgical Gravity-Well Patch into a Standard BF16 Hugging Face Checkpoint | |
| ========================================================================================== | |
| Creates a standard Hugging Face / vLLM / SGLang / LMSYS Arena-compatible BF16 checkpoint | |
| by applying `juker_31b_healed_patch.safetensors` onto `google/gemma-4-31B-it`. | |
| Usage: | |
| python3 apply_patch.py --base-model google/gemma-4-31B-it \ | |
| --patch juker_31b_healed_patch.safetensors \ | |
| --output-dir ./juker-31b | |
| """ | |
| import argparse | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor | |
| from safetensors.torch import load_file | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Merge Juker-31B patch for Hugging Face / Arena deployment") | |
| parser.add_argument("--base-model", type=str, default="google/gemma-4-31B-it", help="Base HF model ID or local path") | |
| parser.add_argument("--patch", type=str, default="juker_31b_healed_patch.safetensors", help="Path to Juker healed patch") | |
| parser.add_argument("--output-dir", type=str, default="./juker-31b", help="Output directory for merged Juker HF model") | |
| parser.add_argument("--push-to-hub", type=str, default=None, help="Optional Hugging Face repo ID to push directly (e.g. username/juker-31b)") | |
| args = parser.parse_args() | |
| print(f"[1/4] Loading base model: {args.base_model} in BF16...") | |
| tokenizer = AutoTokenizer.from_pretrained(args.base_model) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| args.base_model, | |
| torch_dtype=torch.bfloat16, | |
| device_map="cpu", | |
| low_cpu_mem_usage=True, | |
| ) | |
| print(f"[2/4] Applying surgical gravity-well & phase-aligned patch: {args.patch}...") | |
| patch = load_file(args.patch) | |
| layers = model.model.language_model.layers | |
| Q_safety = patch["safety_subspace.Q_safety"].float() | |
| with torch.no_grad(): | |
| for gli in [54, 56, 57, 58, 59]: | |
| cols = patch[f"layers.{gli}.cols"].long() | |
| c_wells = patch[f"layers.{gli}.c_wells"].long() | |
| patched_down = patch[f"layers.{gli}.patched_columns"].to(torch.bfloat16) | |
| patched_gate = patch[f"layers.{gli}.gate_rows"].to(torch.bfloat16) | |
| patched_up = patch[f"layers.{gli}.up_rows"].to(torch.bfloat16) | |
| layers[gli].mlp.down_proj.weight[:, cols] = patched_down | |
| layers[gli].mlp.gate_proj.weight[c_wells, :] = patched_gate | |
| layers[gli].mlp.up_proj.weight[c_wells, :] = patched_up | |
| print(f" -> Patched Layer {gli}: down_proj cols={tuple(patched_down.shape)}, gate/up rows={tuple(patched_gate.shape)}") | |
| print(f"[3/4] Saving standard BF16 Hugging Face / vLLM / SGLang checkpoint to: {args.output_dir}...") | |
| model.save_pretrained(args.output_dir, safe_serialization=True, max_shard_size="5GB") | |
| tokenizer.save_pretrained(args.output_dir) | |
| try: | |
| processor = AutoProcessor.from_pretrained(args.base_model) | |
| processor.save_pretrained(args.output_dir) | |
| except Exception: | |
| pass | |
| if args.push_to_hub: | |
| print(f"[4/4] Pushing merged model to Hugging Face Hub: {args.push_to_hub}...") | |
| model.push_to_hub(args.push_to_hub, safe_serialization=True, max_shard_size="5GB") | |
| tokenizer.push_to_hub(args.push_to_hub) | |
| print("Done! Model is ready for Hugging Face Hub, vLLM, SGLang, and Chatbot Arena.") | |
| if __name__ == "__main__": | |
| main() | |