"""Merge a LoRA adapter into Qwen3.5-9B and save a standalone bf16 model.

    python merge.py runs/run-1/adapters/best /workspace/models/my-model-merged
"""
import sys
import warnings

import torch
from peft import PeftModel
from transformers import (AutoProcessor, AutoTokenizer,
                          Qwen3_5ForConditionalGeneration)

MODEL = "Qwen/Qwen3.5-9B"
REVISION = "c202236235762e1c871ad0ccb60c8ee5ba337b9a"
# Training loaded the text-only model; the full model nests its layers.
KEY_MAPPING = {r"^model\.layers\.": "model.language_model.layers."}

adapter_dir, out = sys.argv[1], sys.argv[2]
warnings.filterwarnings("error", message="Found missing adapter keys")

base = Qwen3_5ForConditionalGeneration.from_pretrained(
    MODEL, revision=REVISION, dtype=torch.bfloat16,
    device_map="cuda:0", attn_implementation="sdpa")
adapter = PeftModel.from_pretrained(base, adapter_dir, key_mapping=KEY_MAPPING)
merged = adapter.merge_and_unload().to(torch.bfloat16)
merged.save_pretrained(out, max_shard_size="4GB")
AutoTokenizer.from_pretrained(MODEL, revision=REVISION).save_pretrained(out)
AutoProcessor.from_pretrained(MODEL, revision=REVISION).save_pretrained(out)
print("merged into", out)