#!/usr/bin/env python3 """Merge LoRA adapter → HuggingFace model → GGUF (blitzkode.gguf). Pipeline -------- 1. Load base model + LoRA adapter, merge and unload adapters. 2. Save merged HuggingFace model to /merged/. 3. Download llama.cpp convert_hf_to_gguf.py from GitHub if not present. 4. Run GGUF conversion (quantised Q4_K_M by default). 5. Verify the output GGUF is loadable with llama-cpp-python. """ from __future__ import annotations import argparse import json import subprocess import sys import urllib.request from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[1] DEFAULT_CHECKPOINT = REPO_ROOT / "checkpoints" / "blitzkode-1.5b-lora" / "final" DEFAULT_MERGED_DIR = REPO_ROOT / "exported" / "merged" DEFAULT_GGUF_OUT = REPO_ROOT / "blitzkode.gguf" LLAMA_CPP_SCRIPTS_DIR = REPO_ROOT / "llama.cpp" CONVERT_SCRIPT_URL = ( "https://raw.githubusercontent.com/ggerganov/llama.cpp/master/convert_hf_to_gguf.py" ) def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__) p.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT) p.add_argument("--merged-dir", type=Path, default=DEFAULT_MERGED_DIR) p.add_argument("--gguf-out", type=Path, default=DEFAULT_GGUF_OUT) p.add_argument("--quant-type", default="q4_k_m", help="GGUF quantisation type (q4_k_m, q8_0, f16, …)") p.add_argument("--skip-merge", action="store_true", help="Skip merge step; use --merged-dir as-is.") p.add_argument("--skip-gguf", action="store_true", help="Skip GGUF conversion (only merge).") p.add_argument("--verify", action="store_true", default=True, help="Verify GGUF is loadable (default: on).") p.add_argument("--no-verify", dest="verify", action="store_false") return p.parse_args() # ─── Step 1: Merge ──────────────────────────────────────────────────────────── def merge_adapter(checkpoint: Path, merged_dir: Path) -> None: print("\n[1/3] Merging LoRA adapter into base model …") import torch # noqa: PLC0415 from peft import PeftModel # noqa: PLC0415 from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: PLC0415 config_path = checkpoint / "adapter_config.json" if not config_path.exists(): raise SystemExit(f"adapter_config.json not found: {config_path}") with config_path.open() as fh: adapter_config = json.load(fh) base_name = adapter_config["base_model_name_or_path"] print(f" Checkpoint : {checkpoint}") print(f" Base model : {base_name}") dtype = torch.float16 print(" Loading base model …") base = AutoModelForCausalLM.from_pretrained( base_name, dtype=dtype, device_map="cpu", trust_remote_code=True, ) tokenizer = AutoTokenizer.from_pretrained(str(checkpoint), trust_remote_code=True) print(" Loading & merging adapter …") model = PeftModel.from_pretrained(base, str(checkpoint)) model = model.merge_and_unload() merged_dir.mkdir(parents=True, exist_ok=True) print(f" Saving merged model to {merged_dir} …") model.save_pretrained(str(merged_dir)) tokenizer.save_pretrained(str(merged_dir)) print(" Merge complete.") # ─── Step 2: Download convert script ───────────────────────────────────────── def ensure_convert_script() -> Path: LLAMA_CPP_SCRIPTS_DIR.mkdir(parents=True, exist_ok=True) convert_script = LLAMA_CPP_SCRIPTS_DIR / "convert_hf_to_gguf.py" if not convert_script.exists(): print(f"\n Downloading convert_hf_to_gguf.py from llama.cpp GitHub …") try: req = urllib.request.Request(CONVERT_SCRIPT_URL, headers={"User-Agent": "BlitzKode/2.0"}) with urllib.request.urlopen(req, timeout=60) as response: content = response.read() convert_script.write_bytes(content) print(f" Saved to: {convert_script}") except Exception as exc: raise SystemExit(f"Failed to download convert_hf_to_gguf.py: {exc}") from exc else: print(f"\n Using cached: {convert_script}") return convert_script # ─── Step 3: GGUF conversion ────────────────────────────────────────────────── def convert_to_gguf(merged_dir: Path, gguf_out: Path, quant_type: str) -> None: print(f"\n[2/3] Converting to GGUF ({quant_type}) …") convert_script = ensure_convert_script() cmd = [ sys.executable, str(convert_script), str(merged_dir), "--outfile", str(gguf_out), "--outtype", quant_type, ] print(f" Running: {' '.join(cmd)}") result = subprocess.run(cmd, capture_output=False, text=True) if result.returncode != 0: raise SystemExit(f"GGUF conversion failed (exit code {result.returncode}).") print(f" GGUF written: {gguf_out} ({gguf_out.stat().st_size / 1024**3:.2f} GB)") # ─── Step 4: Verify ─────────────────────────────────────────────────────────── def verify_gguf(gguf_path: Path) -> None: print(f"\n[3/3] Verifying GGUF with llama-cpp-python …") try: import llama_cpp # noqa: PLC0415 llm = llama_cpp.Llama( model_path=str(gguf_path), n_ctx=128, n_threads=2, n_gpu_layers=0, verbose=False, ) prompt = "<|im_start|>user\nSay hello.<|im_end|>\n<|im_start|>assistant\n" out = llm(prompt, max_tokens=8, stop=["<|im_end|>"]) text = out["choices"][0]["text"].strip() print(f" Sample output: {text!r}") print(" Verification PASSED.") except Exception as exc: print(f" [WARN] Verification raised: {exc}") print(" GGUF was written; manual verification recommended.") # ─── Main ───────────────────────────────────────────────────────────────────── def main() -> None: args = parse_args() print("=" * 72) print("BLITZKODE PRODUCTION EXPORT") print("=" * 72) if not args.skip_merge: merge_adapter(args.checkpoint, args.merged_dir) else: print("\n[1/3] Merge skipped (--skip-merge).") if not args.merged_dir.exists(): raise SystemExit(f"Merged dir not found: {args.merged_dir}") if not args.skip_gguf: convert_to_gguf(args.merged_dir, args.gguf_out, args.quant_type) else: print("\n[2/3] GGUF conversion skipped (--skip-gguf).") if args.verify and args.gguf_out.exists(): verify_gguf(args.gguf_out) print("\n" + "=" * 72) print("EXPORT COMPLETE") print(f" Merged HF model : {args.merged_dir}") if args.gguf_out.exists(): print(f" GGUF model : {args.gguf_out}") print("\nNext steps:") print(" python scripts/push_to_hub.py") print(" python server.py") if __name__ == "__main__": main()