File size: 1,785 Bytes
fedef4a 3d891e5 e824fa1 d884643 fedef4a d884643 e824fa1 fedef4a e824fa1 2292fb9 e824fa1 fedef4a e824fa1 fedef4a e824fa1 fedef4a e824fa1 fedef4a e824fa1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | 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"))
|