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

# %% auto 0
__all__ = ['log', 'prepare_config_for_fp', 'preprocess_config', 'make_datacollator', 'prepare_model_for_training',
           'calculate_token_counts', 'fine_tune', 'prepare_model_for_inference', 'make_pipeline', 'flat_pipeline',
           'generate', 'predict', 'evaluate_']

# %% ../../../nbs/hf.transformers.experiment.ipynb 3
from copy import deepcopy
from math import ceil
from typing import Any, Callable

import torch
from tqdm import tqdm
from datasets import Dataset
from transformers import TrainingArguments, pipeline
from trl import DataCollatorForCompletionOnlyLM, SFTTrainer

from .utils import load_tokenizer_model
from ..datasets.utils import load_datasets
from ...lang.dataset import partition_input_output_messages
from ...logging import get_logger
from ...utils import NestedDict, generate_time_id

log = get_logger(__name__)

# %% ../../../nbs/hf.transformers.experiment.ipynb 4
def prepare_config_for_fp(config: NestedDict):
    if not torch.cuda.is_available():
        return config

    # Set float precision
    if config.at("pretrained_model.torch_dtype") in {"float16", "bfloat16"}:
        major, _ = torch.cuda.get_device_capability()
        gpu_supports_bf = major >= 8
        if gpu_supports_bf:
            log.info("GPU supports bfloat16.")
        else:
            log.info("GPU does not support bfloat16.")
        
        if config.at("pretrained_model.torch_dtype") == "bfloat16" and gpu_supports_bf:
            log.info("Using bfloat16.")
            torch_dtype, bf16, fp16, bnb_4bit_compute_dtype = ("bfloat16", True, False, "bfloat16")
        else:
            log.info("Using float16.")
            torch_dtype, bf16, fp16, bnb_4bit_compute_dtype = ("float16", False, True, "float16")
    else:
        log.info("Not using half-precision float.")
        torch_dtype, bf16, fp16, bnb_4bit_compute_dtype = (None, None, None, None)

    if config.at("pretrained_model.torch_dtype"):
        config.set("pretrained_model.torch_dtype", torch_dtype)
    if config.at("pretrained_model.quantization_config.load_in_4bit"):
        config.set("pretrained_model.quantization_config.bnb_4bit_compute_dtype", bnb_4bit_compute_dtype)
    if config.at("trainer.training_args.bf16") or config.at("trainer.training_args.fp16"):
        config.set("trainer.training_args.bf16", bf16)
        config.set("trainer.training_args.fp16", fp16)
    
    return config

def preprocess_config(config: NestedDict):
    config = deepcopy(config)

    if isinstance(config.at("dataset.train"), dict):
        config.set("dataset.train", [config.at("dataset.train")])
    if isinstance(config.at("dataset.validation"), dict):
        config.set("dataset.validation", [config.at("dataset.validation")])

    config = prepare_config_for_fp(config)
    
    # Generate unique model id for output model
    if (out_model_id := config.at("hfhub.model_id")) and config.at("metaconfig.preprocessing.unique_hfhub_model_id"):
        if "-peft" not in out_model_id and config.at("trainer.lora"):
            out_model_id += "-peft"
        if wandb_run_id := config.at("wandb.id"):
            out_model_id += f"-{wandb_run_id}"
        else:
            out_model_id += f"-{generate_time_id()}"
        config.set("hfhub.model_id", out_model_id)

    return config


# %% ../../../nbs/hf.transformers.experiment.ipynb 6
def make_datacollator(tokenizer, response_template: str | None, response_template_context: str | None = None):
    if not response_template:
        return None

    if response_template_context is None:
        log.info(f"Creating completion-only data collator with response template '{response_template}'")
        data_collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)
    else:
        log.info(f"Creating completion-only data collator with response template '{response_template}' and context '{response_template_context}'")
        response_template_with_context = response_template_context + response_template
        response_template_ids = tokenizer.encode(response_template_with_context, add_special_tokens=False)[len(response_template_context):] 
        data_collator = DataCollatorForCompletionOnlyLM(response_template_ids, tokenizer=tokenizer)
    
    return data_collator

# %% ../../../nbs/hf.transformers.experiment.ipynb 7
def prepare_model_for_training(tokenizer, model):
    model_id = model.name_or_path.lower()

    if "llama-2" in model_id:
        from bellek.hf.transformers.llama2 import prepare_llama2_for_training

        log.info("Base model is a llama-2 model, preparing it for training.")
        prepare_llama2_for_training(tokenizer, model)
    
    elif "llama-3" in model_id:
        from bellek.hf.transformers.llama3 import prepare_llama3_for_training

        log.info("Base model is a llama-3 model, preparing it for training.")
        prepare_llama3_for_training(tokenizer, model)
    
    else:
        log.warning(f"Base model '{model_id}' is not a llama-2 or llama-3 model, no special preparation is done.")

# %% ../../../nbs/hf.transformers.experiment.ipynb 8
def calculate_token_counts(
    tokenizer,
    dataset: Dataset,
    *,
    messages_field: str = "messages",
):
    if messages_field not in dataset.column_names:
        raise ValueError(
            f"Dataset must have `{messages_field}` columns if `text_field` is not specified."
        )

    text_field = "text"
    dataset = dataset.map(
        lambda example: {
            text_field: tokenizer.apply_chat_template(
                example[messages_field],
                tokenize=False,
                add_generation_prompt=False,
            )
        }
    )

    # Inspect token counts
    tokenized_train_ds = dataset.map(lambda examples: tokenizer(examples[text_field]), batched=True)
    token_counts = [len(input_ids) for input_ids in tokenized_train_ds["input_ids"]]
    log.info(f"Input token counts: min={min(token_counts)}, max={max(token_counts)}")
    return token_counts

# %% ../../../nbs/hf.transformers.experiment.ipynb 9
def fine_tune(config: NestedDict):
    from peft import LoraConfig

    # Base model
    pretrained_model_config = config["pretrained_model"]
    model_id = pretrained_model_config.pop("model_name_or_path")
    tokenizer, base_model = load_tokenizer_model(model_id, **pretrained_model_config)
    log.info(f"Loaded base model {model_id}")

    # Prepare model for training
    prepare_model_for_training(tokenizer, base_model)

    # Train dataset
    train_ds = load_datasets(config.at("dataset.train")).shuffle(seed=config.at("seed"))
    log.info(f"Loaded training dataset with {len(train_ds)} samples.")

    # Inspect token counts
    dataset_text_field = config.at("trainer.dataset_text_field")
    token_counts = calculate_token_counts(tokenizer, train_ds)
    log.info(f"Input token counts: min={min(token_counts)}, max={max(token_counts)}")

    # Supervised fine-tuning
    if config.at("trainer.max_seq_length") is None:
        config.set("trainer.max_seq_length", ceil(max(token_counts) / 8) * 8)
    max_seq_length = config.at("trainer.max_seq_length")
    log.info(f"Setting max_seq_length={max_seq_length}")

    peft_config = LoraConfig(**config.at("trainer.lora", {}))

    packing = config.at("trainer.packing", False)

    data_collator = make_datacollator(
        tokenizer,
        config.at("trainer.response_template"),
        config.at("trainer.response_template_context"),
    )

    training_args = TrainingArguments(
        output_dir="./results",
        **config.at("trainer.training_args"),
    )

    trainer = SFTTrainer(
        tokenizer=tokenizer,
        model=base_model,
        peft_config=peft_config,
        train_dataset=train_ds,
        dataset_text_field=dataset_text_field,
        max_seq_length=max_seq_length,
        packing=packing,
        data_collator=data_collator,
        args=training_args,
    )
    log.info("Starting training...")
    trainer.train()

    # Save trained model
    log.info("Saving model...")
    final_model_id = config.at("hfhub.model_id")
    trainer.model.push_to_hub(final_model_id)
    tokenizer.push_to_hub(final_model_id)
    log.info(f"Uploaded PEFT adapters to HF Hub as {final_model_id}")

    return trainer

# %% ../../../nbs/hf.transformers.experiment.ipynb 10
def prepare_model_for_inference(tokenizer, model):
    model_id = model.name_or_path.lower()

    if "llama-2" in model_id:
        from bellek.hf.transformers.llama2 import prepare_llama2_for_inference

        log.info("Base model is a llama-2 model, preparing it for inference.")
        prepare_llama2_for_inference(tokenizer, model)
    
    elif "llama-3" in model_id:
        from bellek.hf.transformers.llama3 import prepare_llama3_for_inference

        log.info("Base model is a llama-3 model, preparing it for inference.")
        prepare_llama3_for_inference(tokenizer, model)
    
    else:
        log.warning(f"Base model '{model_id}' is not a llama-2 or llama-3 model, no special preparation is done.")

# %% ../../../nbs/hf.transformers.experiment.ipynb 11
def _load_tokenizer_model(config: NestedDict):
    model_id = config.at("hfhub.model_id")
    kwargs = deepcopy(config.get("pretrained_model", {}))
    kwargs.pop("model_name_or_path", None)
    return load_tokenizer_model(model_id, **kwargs)


def make_pipeline(config, tokenizer, model):
    prepare_model_for_inference(tokenizer, model)
    return pipeline(
        task="text-generation",
        model=model,
        tokenizer=tokenizer,
        **config.at("inference.pipeline", {}),
    )


def flat_pipeline(pipe):
    def func(inputs, **kwargs) -> list[str]:
        return [result[0]["generated_text"] for result in tqdm(pipe(inputs, **kwargs))]

    return func


def generate(
    pipe,
    inputs,
    **generation_kwargs,
) -> list[str]:
    # Setup generation parameters
    generation_kwargs["return_full_text"] = False

    terminators = generation_kwargs.pop("terminators", [])
    eos_token_ids = {pipe.tokenizer.eos_token_id}
    for terminator in terminators:
        if isinstance(terminator, int):
            eos_token_ids.add(terminator)
        elif isinstance(terminator, str):
            eos_token_ids.add(pipe.tokenizer.convert_tokens_to_ids(terminator))
        else:
            raise ValueError(f"Invalid terminator token {terminator}.")
    generation_kwargs["eos_token_id"] = sorted(eos_token_ids)

    # Generate text
    log.info(f"Running pipeline on {len(inputs)} samples...")
    return flat_pipeline(pipe)(inputs, **generation_kwargs)

def predict(
    config,
    *,
    tokenizer=None,
    model=None,
):
    # Load validation dataset
    dataset_config = config.at("dataset.validation")
    assert dataset_config, "Validation dataset is not provided!"
    dataset = load_datasets(dataset_config)
    assert len(dataset) > 0, "Validation dataset is empty!"

    # Ensure the dataset has input/output columns
    cols = dataset[0].keys()
    if "input" not in cols:
        if "messages" not in dataset.column_names:
            raise ValueError("Dataset must have `messages` column if `input` column are not provided.")
        dataset = dataset.map(partition_input_output_messages).remove_columns("messages")

    # Prepare text generation pipeline
    if tokenizer is None or model is None:
        tokenizer, model = _load_tokenizer_model(config)

    # Set up pipeline
    generation_params = config.at("inference.generation_params", {})
    if "max_new_tokens" not in generation_params:
        if "output" in dataset.column_names:
            token_counts = calculate_token_counts(tokenizer, dataset, messages_field = "output")
            log.info(f"Output token counts: min={min(token_counts)}, max={max(token_counts)}")
            generation_params["max_new_tokens"] = ceil(max(token_counts) / 8) * 8
        else:
            log.info("max_new_tokens is not set.")

    pipe = make_pipeline(config, tokenizer, model)

    # Generate
    generations = generate(
        pipe,
        dataset['input'],
        **generation_params,
    )

    # Create dataframe 
    dataf = dataset.to_pandas()
    dataf["generation"] = generations
    return dataf


def evaluate_(
    config,
    *,
    tokenizer=None,
    model=None,
    metric_kwargs: dict | None = None,
    output_parse_fn: Callable[[str], Any] | None = None,
):
    import evaluate

    output_parse_fn = output_parse_fn or (lambda x: x)

    dataf = predict(config, tokenizer=tokenizer, model=model)

    # Parse texts
    dataf["prediction"] = dataf["generation"].map(output_parse_fn)
    dataf["reference"] = dataf["output"].map(lambda x: x['content']).map(output_parse_fn)

    # Compute scores
    metric = evaluate.load(config.at("evaluation.metric"))
    metric_kwargs = metric_kwargs or {}
    scores = metric.compute(
        predictions=dataf["prediction"].values,
        references=dataf["reference"].values,
        **metric_kwargs,
    )

    return scores, dataf
