# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/ml.peft.ipynb.

# %% auto 0
__all__ = ['log', 'merge_and_publish']

# %% ../../nbs/ml.peft.ipynb 3
import torch
from ..logging import get_logger

log = get_logger(__name__)

# %% ../../nbs/ml.peft.ipynb 4
def merge_and_publish(model_id: str, torch_dtype=torch.float16, device_map={"": 0}, merged_model_id=None):
    from peft import AutoPeftModelForCausalLM
    from transformers import AutoTokenizer
    log.info(f"Loading model and tokenizer for {model_id}") 
    model = AutoPeftModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        device_map=device_map,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    
    log.info("Merging adapters to model...") 
    model = model.merge_and_unload()

    if merged_model_id is None:
        merged_model_id = f"{model_id}-merged"
    log.info(f"Pushing merged model to HF hub as {merged_model_id}") 
    model.push_to_hub(merged_model_id)
    tokenizer.push_to_hub(merged_model_id)
    return merged_model_id
