PYTHON merge_adapter.py

#!/usr/bin/env python3
"""
Merge QLoRA adapter into base model and convert to GGUF for llama.cpp.

Usage:
  python3 merge_adapter.py
"""

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""

from pathlib import Path
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL_NAME = "/home/vincent/projects/bitcoin-ai/model"
ADAPTER_DIR = Path(__file__).parent.parent / "output" / "dev_ai_model" / "final_adapter"
OUTPUT_DIR = Path(__file__).parent.parent / "output" / "merged_model"


def merge():
    print("Loading base model...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        device_map="cpu",
    )
    
    print("Merging adapter...")
    from peft import PeftModel
    model = PeftModel.from_pretrained(model, str(ADAPTER_DIR))
    model = model.merge_and_unload()
    
    print(f"Saving merged model to {OUTPUT_DIR}...")
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(OUTPUT_DIR))
    
    print("Saving tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    tokenizer.save_pretrained(str(OUTPUT_DIR))
    
    print("Merge complete. Run llama.cpp convert script next.")


if __name__ == "__main__":
    merge()