import os
import torch
import json
import shutil
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from huggingface_hub import HfApi

TOKEN = os.environ.get("HF_TOKEN")
api = HfApi(token=TOKEN)

def merge_and_upload(base_id, adapter_id):
    print(f"Merging {adapter_id}...")
    try:
        # Use CPU for merge to avoid OOM if GPU memory is tight
        base = AutoModelForCausalLM.from_pretrained(base_id, torch_dtype=torch.float16, device_map="cpu", token=TOKEN)
        model = PeftModel.from_pretrained(base, adapter_id, token=TOKEN)
        merged = model.merge_and_unload()
        
        out = f"merged_{adapter_id.split('/')[-1]}"
        merged.save_pretrained(out, safe_serialization=True)
        AutoTokenizer.from_pretrained(base_id, token=TOKEN).save_pretrained(out)
        
        # Fix config
        with open(f"{out}/config.json", "r") as f:
            cfg = json.load(f)
        cfg.update({"model_type": "llama", "architectures": ["LlamaForCausalLM"]})
        if "auto_map" in cfg: del cfg["auto_map"]
        with open(f"{out}/config.json", "w") as f:
            json.dump(cfg, f, indent=2)
            
        api.upload_folder(folder_path=out, repo_id=adapter_id, token=TOKEN)
        try:
            api.delete_file("adapter_model.safetensors", repo_id=adapter_id, token=TOKEN)
            api.delete_file("adapter_config.json", repo_id=adapter_id, token=TOKEN)
        except: pass
        shutil.rmtree(out)
        return f"Done {adapter_id}"
    except Exception as e:
        return f"Error {adapter_id}: {str(e)}"

if __name__ == "__main__":
    print(merge_and_upload("NousResearch/Llama-2-7b-hf", "chatpbc1/chatpbc-v4"))
    print(merge_and_upload("NousResearch/Llama-2-7b-hf", "chatpbc1/chatpbc-v33"))