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

# %% auto 0
__all__ = ['log', 'prepare_config_for_fp', 'preprocess_config', 'make_datacollator', 'fine_tune', 'make_io_dataset',
           'make_pipeline', 'flat_pipeline', 'evaluate_pipeline', 'evaluate_']

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

import torch
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 ...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.run_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 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}")

    if "llama" in model_id.lower():
        from bellek.hf.transformers.llama import prepare_llama2_for_training

        log.info("Base model is a LLAMA model, preparing it for training.")
        prepare_llama2_for_training(tokenizer, base_model)

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

    # Convert chat to text
    if "text" not in train_ds.column_names:
        if "messages" not in train_ds.column_names:
            raise ValueError("Dataset does not have 'text' or 'messages' columns.")
        train_ds = train_ds.map(
            lambda example: {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False, add_generation_prompt=False)}
        )

    # Inspect token counts
    tokenized_train_ds = train_ds.map(lambda examples: tokenizer(examples["text"]), 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)}")

    # 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}")

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

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

    peft_config = LoraConfig(**config.at("trainer.lora", {}))
    training_args = TrainingArguments(
        output_dir="./results",
        **config.at("trainer.training_args"),
    )
    trainer = SFTTrainer(
        model=base_model,
        tokenizer=tokenizer,
        peft_config=peft_config,
        train_dataset=train_ds,
        dataset_text_field="text",
        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 8
def make_io_dataset(dataset: Dataset, response_template: str) -> Dataset:
    def extract_input_output(example):
        input, output = example["text"].rsplit(response_template, 1)
        input += response_template
        return {"input": input, "output": output}

    return dataset.map(extract_input_output)


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):
    model_id = config.at("hfhub.model_id")
    if "llama" in model_id:
        from bellek.hf.transformers.llama import prepare_llama2_for_inference

        prepare_llama2_for_inference(tokenizer, model)

    # Create pipeline
    return pipeline(
        task="text-generation",
        model=model,
        tokenizer=tokenizer,
        **config.at("evaluation.pipeline", {}),
    )


def flat_pipeline(pipe):
    def func(inputs) -> list[str]:
        results = pipe(inputs)
        return [result[0]["generated_text"] for result in results]

    return func


def evaluate_pipeline(
    dataset,
    pipe,
    *,
    metric,
    metric_kwargs: dict | None = None,
    output_parse_fn: Callable[[str], Any] | None = None,
):
    eos_token = pipe.tokenizer.special_tokens_map["eos_token"]
    def parse_output(text):
        text = text.replace(eos_token, "").strip()
        if output_parse_fn:
            text = output_parse_fn(text)
        return text

    log.info(f"Running pipeline on dataset with {len(dataset)} samples...")
    generations = flat_pipeline(pipe)(dataset["input"])

    predictions = [parse_output(text) for text in generations]
    references = [parse_output(text) for text in dataset["output"]]

    dataf = dataset.to_pandas()
    dataf["generation"] = generations
    dataf["prediction"] = predictions
    dataf["reference"] = references
    
    metric_kwargs = metric_kwargs or {}
    scores = metric.compute(predictions=predictions, references=references, **metric_kwargs)

    return scores, dataf


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

    # Load validation dataset
    ds_config = config.at("dataset.validation")
    assert ds_config
    ds = load_datasets(ds_config)
    assert len(ds) > 0, "Dataset is empty!"
    
    # Prepare text generation pipeline
    if tokenizer is None or model is None:
        tokenizer, model = _load_tokenizer_model(config)

    # Convert chat to text
    if "text" not in ds.column_names:
        if "messages" not in ds.column_names:
            raise ValueError("Dataset does not have 'text' or 'messages' columns.")
        ds = ds.map(
            lambda example: {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False, add_generation_prompt=False)}
        )

    # Ensure the dataset has input/output columns
    cols = ds[0].keys()
    if "input" not in cols or "output" not in cols:
        response_template = config.at("trainer.response_template")
        assert response_template
        ds = make_io_dataset(ds, response_template)

    if config.at("evaluation.pipeline.max_new_tokens") is None:
        tokenized_outputs = ds.map(lambda examples: tokenizer(examples["output"]), batched=True)
        token_counts = [len(input_ids) for input_ids in tokenized_outputs["input_ids"]]
        config.set("evaluation.pipeline.max_new_tokens", ceil(max(token_counts) / 8) * 8)
        
    log.info(f"Input token counts: min={min(token_counts)}, max={max(token_counts)}")
    pipe = make_pipeline(config, tokenizer, model)

    # Load evaluation metric
    metric = evaluate.load(config.at("evaluation.metric"))

    return evaluate_pipeline(
        ds,
        pipe,
        metric=metric,
        metric_kwargs=metric_kwargs,
        output_parse_fn=output_parse_fn,
    )


