# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/hf.transformers.utils.ipynb.

# %% auto 0
__all__ = ['log', 'prepare_model_kwargs', 'load_tokenizer_model', 'merge_adapters_and_publish']

# %% ../../../nbs/hf.transformers.utils.ipynb 3
from copy import deepcopy
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers.utils import is_flash_attn_2_available
from ...logging import get_logger

log = get_logger(__name__)

# %% ../../../nbs/hf.transformers.utils.ipynb 4
def prepare_model_kwargs(model_kwargs: dict) -> dict:
    model_kwargs = deepcopy(model_kwargs)
    
    # Setup quantization config 
    if (quantization_config := model_kwargs.get("quantization_config")) and isinstance(quantization_config, dict):
        from transformers import BitsAndBytesConfig

        model_kwargs["quantization_config"] = BitsAndBytesConfig(**quantization_config)
    
    # Setup torch dtype
    if (torch_dtype := model_kwargs.get("torch_dtype")) and (torch_dtype != "auto"):
        model_kwargs["torch_dtype"] = getattr(torch, torch_dtype)

    # Setup device map
    if "device_map" not in model_kwargs:
        model_kwargs["device_map"] = {"": torch.cuda.current_device()}

    # Setup flash attention 2
    if model_kwargs.get("attn_implementation") == "flash_attention_2":
        if not is_flash_attn_2_available():
            log.warning("Flash attention 2 is not available. Falling back to default attention.")
            del model_kwargs["attn_implementation"]
    return model_kwargs

# %% ../../../nbs/hf.transformers.utils.ipynb 5
def load_tokenizer_model(
    model_name_or_path: str,
    *,
    auto_model_cls=None,
    **model_kwargs,
):

    model_kwargs = prepare_model_kwargs(model_kwargs)

    if auto_model_cls is None:
        if "-peft" in model_name_or_path:
            from peft import AutoPeftModelForCausalLM

            auto_model_cls = AutoPeftModelForCausalLM
        else:
            auto_model_cls = AutoModelForCausalLM

    # Load model
    model = auto_model_cls.from_pretrained(
        model_name_or_path,
        **model_kwargs,
    )
    # Load tokenizer
    if auto_model_cls == AutoModelForCausalLM:
        tokenizer_id = model_name_or_path
    else:
        from peft import AutoPeftModelForCausalLM

        if auto_model_cls == AutoPeftModelForCausalLM:
            tokenizer_id = model.active_peft_config.base_model_name_or_path
        else:
            raise ValueError(f"Unknown auto_model_cls: {auto_model_cls}")
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, trust_remote_code=True)
    return tokenizer, model

# %% ../../../nbs/hf.transformers.utils.ipynb 6
def merge_adapters_and_publish(
    model_id: str,
    torch_dtype=torch.float16,
    device_map={"": 0},
    merged_model_id: str=None,
):
    from peft import AutoPeftModelForCausalLM
    
    if isinstance(torch_dtype, str) and torch_dtype != "auto":
        torch_dtype = getattr(torch, torch_dtype) 

    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
