# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" Training DPO """

import os
import sys
import time
from functools import partial

import paddle

from llm.alignment.dpo.dpo_argument import (
    DPOConfig,
    DPODataArgument,
    DPOModelArgument,
    DPOTrainingArguments,
)
from paddlenlp.datasets import (
    ZeroPaddingIterableDataset,
    ZeroPaddingMapDataset,
    load_dataset,
)
from paddlenlp.peft import LoRAConfig, LoRAModel
from paddlenlp.trainer import PdArgumentParser, get_last_checkpoint, set_seed
from paddlenlp.transformers import (
    AutoConfig,
    AutoModelForCausalLMPipe,
    AutoTokenizer,
    LlamaForCausalLMNetDPO,
)
from paddlenlp.transformers.configuration_utils import LlmMetaConfig
from paddlenlp.trl import (
    DPOAutoTrainer,
    calculate_effective_tokens,
    preference_collate_fn_auto_parallel,
    preprocess_preference_data,
)
from paddlenlp.trl.llm_utils import get_lora_target_modules
from paddlenlp.utils.log import logger

flash_mask_support_list = [
    LlamaForCausalLMNetDPO,
]


def main():
    """main"""
    parser = PdArgumentParser((DPOModelArgument, DPODataArgument, DPOTrainingArguments, DPOConfig))
    if len(sys.argv) >= 2 and sys.argv[1].endswith(".json"):
        model_args, data_args, training_args, dpo_config = parser.parse_json_file_and_cmd_lines()
    else:
        model_args, data_args, training_args, dpo_config = parser.parse_args_into_dataclasses()

    paddle.set_device(training_args.device)
    set_seed(training_args.seed)
    if dpo_config.loss_type == "orpo":
        dpo_config.reference_free = True
        dpo_config.sft_loss_ratio = 1.0
        dpo_config.loss_type = "or"
        logger.info("orpo loss_type is equal to sft_loss + pref_loss_ratio * or_loss.")
    if dpo_config.loss_type in ["or", "simpo"] and not dpo_config.reference_free:
        dpo_config.reference_free = True
        logger.warning(f"{dpo_config.loss_type} loss_type only supports reference_free. Set reference_free to True.")
    if training_args.pipeline_parallel_degree > 1:
        assert (
            hasattr(training_args, "pipeline_parallel_config")
            and "enable_clear_every_step_cache" in training_args.pipeline_parallel_config
        ), "Should set '--pipeline_parallel_config enable_clear_every_step_cache' in bash script for pp."
    if training_args.sequence_parallel:
        if training_args.pipeline_parallel_degree > 1:
            assert (
                hasattr(training_args, "pipeline_parallel_config")
                and "disable_partial_send_recv" in training_args.pipeline_parallel_config
            ), "Should set '--pipeline_parallel_config disable_partial_send_recv' in bash script for pp with sp."
        if training_args.tensor_parallel_degree <= 1:
            training_args.sequence_parallel = False
            logger.info("Tensor_parallel_degree = 1. Set sequence_parallel to False.")
    training_args.print_config(model_args, "Model")
    training_args.print_config(data_args, "Data")
    training_args.print_config(dpo_config, "DPOConfig")

    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, world_size: "
        f"{training_args.world_size}, distributed training: {bool(training_args.local_rank != -1)}, "
        f"16-bits training: {training_args.fp16 or training_args.bf16}"
    )

    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set the dtype for loading model
    dtype = paddle.get_default_dtype()
    if training_args.fp16_opt_level == "O2":
        if training_args.fp16:
            dtype = "float16"
        if training_args.bf16:
            dtype = "bfloat16"

    logger.info("Start to load model & tokenizer.")

    model_config = AutoConfig.from_pretrained(model_args.model_name_or_path, dtype=dtype)
    model_config["num_hidden_layers"] = training_args.num_hidden_layers
    LlmMetaConfig.set_llm_config(model_config, training_args)

    if not dpo_config.reference_free and not dpo_config.lora:
        ref_model_config = AutoConfig.from_pretrained(model_args.model_name_or_path, dtype=dtype)
        ref_model_config["num_hidden_layers"] = training_args.num_hidden_layers
        LlmMetaConfig.set_llm_config(ref_model_config, training_args)

    if training_args.pipeline_parallel_degree > 1:
        model_class = AutoModelForCausalLMPipe
        if not dpo_config.reference_free and not dpo_config.lora:
            ref_model_config.dpo_config = dpo_config
        model_config.dpo_config = dpo_config
    else:
        model_class = LlamaForCausalLMNetDPO

    if not training_args.autotuner_benchmark or model_args.weight_quantize_algo is not None:
        model = model_class.from_pretrained(model_args.model_name_or_path, config=model_config)
        # for DPO save
        if not dpo_config.reference_free and not dpo_config.lora:
            model_config["num_hidden_layers"] = training_args.num_hidden_layers
            ref_model_config["num_hidden_layers"] = training_args.num_hidden_layers
            ref_model = model_class.from_config(ref_model_config)
            ref_model.set_state_dict(model.state_dict())
        else:
            ref_model = None
    else:
        model = model_class.from_config(model_config)
        if not dpo_config.reference_free and not dpo_config.lora:
            ref_model = model_class.from_config(ref_model_config)
        else:
            ref_model = None
    if training_args.pipeline_parallel_degree > 1:
        model.config.dpo_config = None
    if model_args.flash_mask and not model.config.use_flash_attention:
        logger.warning("`flash_mask` must use with zero padding and flash attention.")
        model.config.use_flash_attention = True

    if model_args.flash_mask and not any(isinstance(model, cls) for cls in flash_mask_support_list):
        raise NotImplementedError(f"{model.__class__} not support flash mask.")

    if model_args.tokenizer_name_or_path is not None:
        tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path)
    else:
        tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
    # TODO: support chat template in next pr
    tokenizer.chat_template = None
    logger.info("Loading model & tokenizer successfully !")

    if dpo_config.lora:
        if training_args.sharding_parallel_degree > 1:
            assert (
                "enable_stage1_overlap" not in training_args.sharding_parallel_config
            ), "Currently not support enabling sharding_stage1_overlap in lora mode."
        if model_args.lora_path is None:
            target_modules = get_lora_target_modules(model)
            if model_args.rslora_plus:
                model_args.rslora = True
                model_args.lora_plus_scale = 4
                model_args.lora_alpha = 4
            if model_args.weight_quantize_algo is not None:
                if model_args.rslora or model_args.lora_plus_scale != 1.0:
                    logger.info("Weight quantization is not supported in LoRA+ and RsLoRA.")
            if model_args.lora_alpha == -1:
                if model_args.rslora:
                    model_args.lora_alpha = 4
                else:
                    model_args.lora_alpha = 2 * model_args.lora_rank
            lora_config = LoRAConfig(
                target_modules=target_modules,
                r=model_args.lora_rank,
                lora_alpha=2 * model_args.lora_rank if not model_args.rslora else 4,
                rslora=model_args.rslora,
                lora_plus_scale=model_args.lora_plus_scale,
                tensor_parallel_degree=training_args.tensor_parallel_degree,
                dtype=dtype,
                base_model_name_or_path=model_args.model_name_or_path,
                use_quick_lora=model_args.use_quick_lora,
            )
            model = LoRAModel(model, lora_config)
        else:
            model = LoRAModel.from_pretrained(model=model, lora_path=model_args.lora_path)

        model.print_trainable_parameters()

    logger.info("Start to create dataset")
    trans_func = partial(preprocess_preference_data, tokenizer=tokenizer, data_args=data_args, model_args=model_args)
    if data_args.lazy:
        zero_padding_dataset = ZeroPaddingIterableDataset
    else:
        zero_padding_dataset = ZeroPaddingMapDataset
    if training_args.do_train and training_args.should_load_dataset:
        train_ds = load_dataset(
            "json",
            data_files=data_args.train_dataset_path,
            lazy=data_args.lazy,
        )[0]
        logger.info("Creating train Zero Padding Data Stream. This may take a few minutes.")
        train_ds = (
            zero_padding_dataset(
                train_ds.map(trans_func),
                tokenizer=tokenizer,
                max_length=data_args.max_seq_len,
                greedy_zero_padding=data_args.greedy_zero_padding,
            )
            if train_ds is not None
            else None
        )
    else:
        train_ds = None

    if training_args.do_eval and training_args.should_load_dataset:
        eval_ds = load_dataset(
            "json",
            data_files=data_args.dev_dataset_path,
            lazy=data_args.lazy,
        )[0]
        logger.info("Creating dev Zero Padding Data Stream. This may take a few minutes.")
        eval_ds = (
            zero_padding_dataset(
                eval_ds.map(trans_func),
                tokenizer=tokenizer,
                max_length=data_args.max_seq_len,
            )
            if eval_ds is not None
            else None
        )
    else:
        eval_ds = None
    logger.info("Creating dataset successfully ...")

    trainer = DPOAutoTrainer(
        model=model,
        ref_model=ref_model,
        dpo_config=dpo_config,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        tokenizer=tokenizer,
        data_collator=partial(
            preference_collate_fn_auto_parallel,
            max_seq_len=data_args.max_seq_len,
        ),
        ignore_eos_token=True,
    )

    if training_args.do_train:
        train_result = trainer.train(resume_from_checkpoint=last_checkpoint)

        if training_args.benchmark:
            total_effective_tokens, total_tokens = calculate_effective_tokens(
                training_args, train_ds, data_args.max_seq_len
            )
            effective_tokens_per_second = total_effective_tokens / train_result.metrics["train_runtime"]
            total_tokens_per_second = total_tokens / train_result.metrics["train_runtime"]
            effective_ratio = 100 * total_effective_tokens / total_tokens
            logger.info(
                "[timelog] {}: {:.2f} % ({}) ".format(
                    "Effective ratio", effective_ratio, time.strftime("%Y-%m-%d %H:%M:%S")
                )
            )
            logger.info(
                "[timelog] {}: {:.2f} token/s ({}) ".format(
                    "Effective tokens per second", effective_tokens_per_second, time.strftime("%Y-%m-%d %H:%M:%S")
                )
            )
            logger.info(
                "[timelog] {}: {:.2f} token/s ({}) ".format(
                    "Tokens per second", total_tokens_per_second, time.strftime("%Y-%m-%d %H:%M:%S")
                )
            )


if __name__ == "__main__":
    main()
