| |
| """ |
| 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() |
|
|
