68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Merge a LoRA adapter (saved by your finetune run) into the full-precision
|
|
base model and write the merged fp16 weights to disk.
|
|
|
|
Requires:
|
|
pip install unsloth peft transformers accelerate
|
|
"""
|
|
|
|
from pathlib import Path
|
|
import torch
|
|
from unsloth import FastLanguageModel
|
|
from peft import PeftModel
|
|
from transformers import AutoTokenizer
|
|
|
|
# ---------------------------------------------------------------------
|
|
# Adjust these three paths if your directory layout is different.
|
|
# ---------------------------------------------------------------------
|
|
BASE_MODEL = "unsloth/gemma-3-4b-it"
|
|
ADAPTER_DIR = (
|
|
"results/auto_antislop_runs/run_20250608_102159/"
|
|
"finetuned_model_ftpo_exp01/lora_adapters"
|
|
)
|
|
OUT_DIR = (
|
|
"results/auto_antislop_runs/run_20250608_102159/"
|
|
"finetuned_model_ftpo_exp01/merged_manual_fp16"
|
|
)
|
|
|
|
# ---------------------------------------------------------------------
|
|
def main() -> None:
|
|
print("→ loading base model …")
|
|
base_model, _ = FastLanguageModel.from_pretrained(
|
|
model_name = BASE_MODEL,
|
|
max_seq_length = 4096, # keep consistent with training
|
|
load_in_4bit = False, # full-precision
|
|
dtype = torch.float16,
|
|
device_map = {"": "cpu"}, # CPU merge; change to {"": 0} for GPU
|
|
)
|
|
|
|
print("→ plugging in LoRA adapter …")
|
|
peft_model = PeftModel.from_pretrained(
|
|
base_model,
|
|
ADAPTER_DIR,
|
|
device_map = {"": "cpu"},
|
|
)
|
|
|
|
print("→ merging and unloading …")
|
|
merged_model = peft_model.merge_and_unload() # returns a plain nn.Module
|
|
|
|
Path(OUT_DIR).mkdir(parents=True, exist_ok=True)
|
|
|
|
print(f"→ saving merged model to {OUT_DIR} …")
|
|
merged_model.save_pretrained(
|
|
OUT_DIR,
|
|
safe_serialization = True, # *.safetensors shards
|
|
max_shard_size = "5GB",
|
|
)
|
|
|
|
# save the tokenizer so the directory is immediately usable
|
|
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
|
|
tokenizer.save_pretrained(OUT_DIR)
|
|
|
|
print("✓ done")
|
|
|
|
# ---------------------------------------------------------------------
|
|
if __name__ == "__main__":
|
|
main()
|