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

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

# %% ../../nbs/ml.transformers.ipynb 3
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from ..logging import get_logger

log = get_logger(__name__)

# %% ../../nbs/ml.transformers.ipynb 4
def merge_adapters_and_publish(
    model_id: str,
    torch_dtype=torch.float16,
    device_map={"": 0},
    merged_model_id: str=None,
):
    from peft import AutoPeftModelForCausalLM

    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

# %% ../../nbs/ml.transformers.ipynb 5
def load_tokenizer_model(
    model_name_or_path: str,
    auto_model_cls=AutoModelForCausalLM,
    quantization_config=None,
    device_map={"": 0},
):
    # Setup quantization config
    if isinstance(quantization_config, dict):
        from transformers import BitsAndBytesConfig
        quantization_config = BitsAndBytesConfig(**quantization_config)
    
    # Load model
    model = auto_model_cls.from_pretrained(
        model_name_or_path,
        device_map=device_map,
        quantization_config=quantization_config,
    )

    # 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
