# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by the HuggingFace's TRL library.
# https://github.com/huggingface/trl/blob/v0.8.0/trl/trainer/dpo_trainer.py
#
# 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.

import numpy as np
import warnings
from collections import defaultdict
from contextlib import nullcontext
from types import MethodType
from typing import TYPE_CHECKING, Dict, Literal, Optional, Tuple, Union

import torch
import os
import torch.distributed as dist
import torch.nn.functional as F
from torch.utils.data import SequentialSampler
from transformers import Trainer
from trl import DPOTrainer
from trl.trainer import disable_dropout_in_model
from typing_extensions import override

from ...extras.constants import IGNORE_INDEX
from ...extras.packages import is_transformers_version_equal_to_4_46
from ..callbacks import PissaConvertCallback, SaveProcessorCallback
from ..trainer_utils import create_custom_optimizer, create_custom_scheduler, get_batch_logps


if TYPE_CHECKING:
    from transformers import PreTrainedModel, ProcessorMixin

    from ...hparams import FinetuningArguments


class CustomDPOTrainer(DPOTrainer):
    def __init__(
        self,
        model: Union["PreTrainedModel", torch.nn.Module],
        ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]],
        finetuning_args: "FinetuningArguments",
        processor: Optional["ProcessorMixin"],
        disable_dropout: bool = True,
        **kwargs,
    ):
        if disable_dropout:
            disable_dropout_in_model(model)
            if ref_model is not None:
                disable_dropout_in_model(ref_model)

        self.finetuning_args = finetuning_args
        self.f_divergence_type = "reverse_kl"
        self.reference_free = False
        self.use_dpo_data_collator = True  # hack to avoid warning
        self.generate_during_eval = False  # disable at evaluation
        self.label_pad_token_id = IGNORE_INDEX
        self.padding_value = 0
        self.is_encoder_decoder = model.config.is_encoder_decoder
        self.precompute_ref_log_probs = False
        self._precomputed_train_ref_log_probs = False
        self._precomputed_eval_ref_log_probs = False
        self._peft_has_been_casted_to_bf16 = False

        self.ref_model = ref_model
        self._stored_metrics = defaultdict(lambda: defaultdict(list))

        self._has_dummy_forwarded = False

        # shor2long
        self.add_short2long_loss = finetuning_args.add_short2long_loss
        self.kl_lambda = finetuning_args.kl_lambda
        self.kl_penalty = finetuning_args.kl_penalty
        self.kl_part = finetuning_args.kl_part
        self.kl_reverse = finetuning_args.kl_reverse
        self.efficiency = finetuning_args.efficiency
        self.seq_forward = finetuning_args.seq_forward
        print("==="*5,"short2long","==="*5)
        print(f"add_short2long_loss: {self.add_short2long_loss}")
        # if self.loss_type == "orpo":
        #     self.kl_lambda = self.kl_lambda * self.beta
        #     print("self.kl_lambda = self.kl_lambda * self.beta")
        print(f"seq_forward: {self.seq_forward}")
        print(f"efficiency: {self.efficiency}")
        print(f"kl_lambda: {self.kl_lambda}")
        print(f"kl_part: {self.kl_part}")
        print(f"kl_penalty: {self.kl_penalty}")
        print(f"kl_reverse: {self.kl_reverse}")
        print("==="*5,"short2long","==="*5)

        # dpo hyperparams
        self.beta = finetuning_args.pref_beta
        self.loss_type = finetuning_args.pref_loss
        self.ftx_gamma = finetuning_args.pref_ftx
        self.label_smoothing = finetuning_args.dpo_label_smoothing
        self.simpo_gamma = finetuning_args.simpo_gamma
        self.beta_sla = finetuning_args.beta_sla
        self.eta_sla = finetuning_args.eta_sla
        self.gamma_sla = finetuning_args.gamma_sla
        self.sft_part = finetuning_args.sft_part

        Trainer.__init__(self, model=model, **kwargs)
        if not hasattr(self, "accelerator"):
            raise AttributeError("Please update `transformers`.")

        warnings.simplefilter("ignore")  # remove gc warnings on ref model

        if ref_model is not None:
            if self.is_deepspeed_enabled:
                if not (
                    getattr(ref_model, "is_loaded_in_8bit", False) or getattr(ref_model, "is_loaded_in_4bit", False)
                ):  # quantized models are already set on the correct device
                    self.ref_model = self._prepare_deepspeed(self.ref_model)
            else:
                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
                self.ref_model.eval()

        if processor is not None:
            self.add_callback(SaveProcessorCallback(processor))

        if finetuning_args.pissa_convert:
            self.callback_handler.add_callback(PissaConvertCallback)

        if finetuning_args.use_badam:
            from badam import BAdamCallback, clip_grad_norm_old_version  # type: ignore

            self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
            self.add_callback(BAdamCallback)

    @override
    def create_optimizer(self) -> "torch.optim.Optimizer":
        if self.optimizer is None:
            self.optimizer = create_custom_optimizer(self.model, self.args, self.finetuning_args)
        return super().create_optimizer()

    @override
    def create_scheduler(
        self, num_training_steps: int, optimizer: Optional["torch.optim.Optimizer"] = None
    ) -> "torch.optim.lr_scheduler.LRScheduler":
        create_custom_scheduler(self.args, num_training_steps, optimizer)
        return super().create_scheduler(num_training_steps, optimizer)

    @override
    def get_batch_samples(self, epoch_iterator, num_batches):
        r"""
        Replaces the method of KTO Trainer with the one of the standard Trainer.
        """
        return Trainer.get_batch_samples(self, epoch_iterator, num_batches)

    def log_odds(self, chosen_logps: "torch.Tensor", rejected_logps: "torch.Tensor") -> "torch.Tensor":
        r"""
        Computes ORPO's odds ratio (OR) loss for batched log probabilities of the policy model.
        """
        log_odds = (chosen_logps - rejected_logps) - (
            torch.log1p(-torch.exp(chosen_logps)) - torch.log1p(-torch.exp(rejected_logps))
        )
        return log_odds

    def odds_ratio_loss(self, chosen_logps: "torch.Tensor", rejected_logps: "torch.Tensor") -> "torch.Tensor":
        r"""
        Computes ORPO's odds ratio (OR) loss for batched log probabilities of the policy model.
        """
        log_odds = (chosen_logps - rejected_logps) - (
            torch.log1p(-torch.exp(chosen_logps)) - torch.log1p(-torch.exp(rejected_logps))
        )
        sft_loss = -chosen_logps
        odds_ratio_loss = -F.logsigmoid(log_odds)
        orpo_loss = sft_loss + self.beta * odds_ratio_loss
        return orpo_loss

    def simpo_loss(self, chosen_logps: "torch.Tensor", rejected_logps: "torch.Tensor",) -> "torch.Tensor":
        r"""
        Computes SimPO loss for batched log probabilities of the policy model.
        """
        pi_logratios = chosen_logps - rejected_logps
        gamma_logratios = self.simpo_gamma / (self.beta * self.beta_sla)
        logits = pi_logratios - gamma_logratios
        simpo_loss = -F.logsigmoid((self.beta * self.beta_sla) * logits)
        return simpo_loss

    def compute_preference_loss(
        self,
        policy_chosen_logps: "torch.Tensor",
        policy_rejected_logps: "torch.Tensor",
        reference_chosen_logps: Optional["torch.Tensor"],
        reference_rejected_logps: Optional["torch.Tensor"],
    ) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
        r"""
        Computes loss for preference learning.
        """
        if self.loss_type in ["simpo", "orpo"]:
            if self.loss_type == "orpo":
                losses = self.odds_ratio_loss(policy_chosen_logps, policy_rejected_logps)
            elif self.loss_type == "simpo":
                losses = self.simpo_loss(policy_chosen_logps, policy_rejected_logps)
            else:
                raise NotImplementedError(f"Unknown loss type: {self.loss_type}.")

            chosen_rewards = self.beta * policy_chosen_logps.to(self.accelerator.device).detach()
            rejected_rewards = self.beta * policy_rejected_logps.to(self.accelerator.device).detach()
        else:
            losses, chosen_rewards, rejected_rewards = self.dpo_loss(
                policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps
            )
        return losses, chosen_rewards, rejected_rewards

    @override
    def concatenated_forward(
        self, model: "PreTrainedModel", batch: Dict[str, "torch.Tensor"]
    ) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]:
        r"""
        Computes the sum log probabilities of the labels under given logits if loss_type is not IPO, ORPO or SimPO.

        Otherwise the average log probabilities.
        """
        # print(batch.keys())
        # ['input_ids', 'attention_mask', 'labels']
        # exit()
        if self.finetuning_args.use_ref_model:
            batch = {k: v.detach().clone() for k, v in batch.items()}  # avoid error

        all_logits: "torch.Tensor" = model(**batch, return_dict=True, use_cache=False).logits.to(torch.float32)
        all_logps, valid_length = get_batch_logps(
            logits=all_logits, labels=batch["labels"], shift_labels=model.sequence_parallel_group is None
        )  # shift labels if no sequence parallel
        if self.loss_type in ["ipo", "orpo", "simpo"]:
            # add for sequence parallel
            sp_group = model.sequence_parallel_group
            if sp_group is not None:
                dist.all_reduce(all_logps, op=dist.ReduceOp.SUM, group=sp_group)
                dist.all_reduce(valid_length, op=dist.ReduceOp.SUM, group=sp_group)
            all_logps = all_logps / valid_length

        batch_size = batch["input_ids"].size(0) // 2
        chosen_logps, rejected_logps = all_logps.split(batch_size, dim=0)
        chosen_logits, rejected_logits = all_logits.split(batch_size, dim=0)
        chosen_length, _ = valid_length.split(batch_size, dim=0)
        return chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_length

    @override
    def compute_reference_log_probs(
        self, model: "PreTrainedModel", batch: Dict[str, "torch.Tensor"]
    ) -> Tuple[Optional["torch.Tensor"], Optional["torch.Tensor"]]:
        r"""
        Computes log probabilities of the reference model.
        """
        if not self.finetuning_args.use_ref_model:
            return None, None

        if self.ref_model is None:
            ref_model = model
            ref_context = self.accelerator.unwrap_model(model).disable_adapter()
        else:
            ref_model = self.ref_model
            ref_context = nullcontext()

        with torch.no_grad(), ref_context:
            reference_chosen_logps, reference_rejected_logps, *_ = self.concatenated_forward(ref_model, batch)

        return reference_chosen_logps, reference_rejected_logps

    def compute_kl_penalty(self, logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty, long_policy_chosen_length) -> torch.FloatTensor:
        """Compute KL divergence given logprob and ref_logprob.
        Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1104

        Args:
            logprob:
            ref_logprob:

        Returns:

        """
        if kl_penalty == "kl":
            return logprob - ref_logprob if not self.kl_reverse else ref_logprob - logprob

        if kl_penalty == "abs":
            return (logprob - ref_logprob).abs()

        if kl_penalty == "mse":
            return 0.5 * (logprob - ref_logprob).square()

        # J. Schulman. Approximating kl divergence, 2020.
        # # URL http://joschu.net/blog/kl-approx.html.
        if kl_penalty == 'low_var_kl':
            kl = ref_logprob - logprob if not self.kl_reverse else logprob - ref_logprob
            ratio = torch.exp(kl)
            kld = (ratio - kl - 1).contiguous()
            return torch.clamp(kld, min=-10, max=10)

        if kl_penalty == 'sft':
            # policy_chosen_logps_avg = policy_chosen_logps / policy_chosen_length
            # sft_loss = -policy_chosen_logps_avg
            return -1 * logprob

        if kl_penalty == "full":
            # so, here logprob and ref_logprob should contain the logits for every token in vocabulary
            raise NotImplementedError

        raise NotImplementedError

    def compute_short2long_loss(
        self,
        policy_chosen_logps,
        policy_rejected_logps,
        long_policy_chosen_logps,
        long_policy_rejected_logps,
        long_policy_chosen_length = None
    ):
        r"""
        Compute the loss for short2long.
        Specifially, the loss is computed by KL divergence.
        args:
            long_chosen_policy_logp: [batch_size]
            short_chosen_policy_logp: [batch_size]
            long_rejected_policy_logp: [batch_size]
            short_rejected_policy_logp: [batch_size]
        """
        # GRPO kl
        # if kl_penalty == 'low_var_kl':
        # kl = ref_logprob - logprob
        # ratio = torch.exp(kl)
        # kld = (ratio - kl - 1).contiguous()
        # return torch.clamp(kld, min=-10, max=10)
        if self.kl_part == "both":
            chosen_logratios = self.compute_kl_penalty(long_policy_chosen_logps, policy_chosen_logps, self.kl_penalty, long_policy_chosen_length)
            rejected_logratios = self.compute_kl_penalty(long_policy_rejected_logps, policy_rejected_logps, self.kl_penalty, long_policy_chosen_length)
            kl_chosen_loss = chosen_logratios.to(self.accelerator.device)
            kl_rejected_loss = rejected_logratios.to(self.accelerator.device)
            return kl_chosen_loss, kl_rejected_loss
        if self.kl_part == "chosen":
            chosen_logratios = self.compute_kl_penalty(long_policy_chosen_logps, policy_chosen_logps, self.kl_penalty, long_policy_chosen_length)
            kl_chosen_loss = chosen_logratios.to(self.accelerator.device)
            return kl_chosen_loss, 0
        if self.kl_part == "rejected":
            rejected_logratios = self.compute_kl_penalty(long_policy_rejected_logps, policy_rejected_logps, self.kl_penalty, long_policy_chosen_length)
            kl_rejected_loss = rejected_logratios.to(self.accelerator.device)
            return 0, kl_rejected_loss
        return 0, 0

    def compute_long_rewards(
        self,
        long_policy_chosen_logps,
        long_policy_rejected_logps,
        policy_chosen_logps,
        policy_rejected_logps,
    ):
        # long_policy_chosen_logps
        # long_policy_rejected_logps
        if self.loss_type == "sigmoid":
            long_dpo_losses, long_chosen_rewards, long_rejected_rewards = self.dpo_loss(
                            long_policy_chosen_logps, long_policy_rejected_logps, policy_chosen_logps, policy_rejected_logps
                        )
        else:
            long_chosen_rewards = self.beta * long_policy_chosen_logps.to(self.accelerator.device).detach()
            long_rejected_rewards = self.beta * long_policy_rejected_logps.to(self.accelerator.device).detach()
        return long_chosen_rewards, long_rejected_rewards
    
    def separate_forward(self, model: "PreTrainedModel", batch: Dict[str, "torch.Tensor"]):
        if self.finetuning_args.use_ref_model:
            batch = {k: v.detach().clone() for k, v in batch.items()}  # avoid error

        all_logits: "torch.Tensor" = model(**batch, return_dict=True, use_cache=False).logits.to(torch.float32)
        all_logps, valid_length = get_batch_logps(
            logits=all_logits, labels=batch["labels"], shift_labels=model.sequence_parallel_group is None
        )  # shift labels if no sequence parallel

        if self.loss_type in ["ipo", "orpo", "simpo"]:
            # add for sequence parallel
            sp_group = model.sequence_parallel_group
            if sp_group is not None:
                dist.all_reduce(all_logps, op=dist.ReduceOp.SUM, group=sp_group)
                dist.all_reduce(valid_length, op=dist.ReduceOp.SUM, group=sp_group)
            all_logps = all_logps / valid_length
        return all_logps, all_logits, valid_length

    def dummy_concatenated_forward(self, model: "PreTrainedModel", batch: Dict[str, "torch.Tensor"], chosen_only=False
    ) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]:
        batch_size = batch["input_ids"].size(0) // 2
        chosen_batch = {}
        rejected_batch = {}
        for k,v in batch.items():
            chosen_split,rejected_split = v.split(batch_size, dim=0)
            chosen_batch[k] = chosen_split
            rejected_batch[k] = rejected_split
        chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_length = None, None, None, None, None
        if chosen_only:
            chosen_logps, chosen_logits, chosen_length = self.separate_forward(model, chosen_batch)
        else:
            chosen_logps, chosen_logits, chosen_length = self.separate_forward(model, chosen_batch)
            rejected_logps, rejected_logits, _ = self.separate_forward(model, rejected_batch)

        return chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_length

    def dummy_compute_reference_log_probs(
        self, model: "PreTrainedModel", batch: Dict[str, "torch.Tensor"], chosen_only=False
    ) -> Tuple[Optional["torch.Tensor"], Optional["torch.Tensor"]]:

        if not self.finetuning_args.use_ref_model:
            return None, None
        if self.ref_model is None:
            ref_model = model
            ref_context = self.accelerator.unwrap_model(model).disable_adapter()
        else:
            ref_model = self.ref_model
            ref_context = nullcontext()
            
        with torch.no_grad(), ref_context:
            reference_chosen_logps, reference_rejected_logps, *_ = self.dummy_concatenated_forward(ref_model, batch, chosen_only=chosen_only)

        return reference_chosen_logps, reference_rejected_logps

    @override
    def get_batch_loss_metrics(
        self,
        model: "PreTrainedModel",
        batch: Dict[str, "torch.Tensor"],
        train_eval: Literal["train", "eval"] = "train",
    ) -> Tuple["torch.Tensor", Dict[str, "torch.Tensor"]]:
        r"""
        Computes the DPO loss and other metrics for the given batch of inputs for train or test.
        """
        # ['input_ids', 'attention_mask', 'labels']
        # current_dir = os.path.dirname(os.path.abspath(__file__))
        # print(batch.keys())
        # print("input_ids.size()", batch["input_ids"].size())
        # [2, 2816]
        # [batch_size*2, sequlence_length]
        # exit()
        long_batch = None
        if "long_input_ids" in batch.keys():
            long_batch = {}
            long_batch["input_ids"] = batch.pop("long_input_ids")
            long_batch["attention_mask"] = batch.pop("long_attention_mask")
            long_batch["labels"] = batch.pop("long_labels")

        metrics = {}
        if not self.seq_forward:
            (
                policy_chosen_logps,
                policy_rejected_logps,
                policy_chosen_logits,
                policy_rejected_logits,
                policy_chosen_length,
            ) = self.concatenated_forward(model, batch)
            reference_chosen_logps, reference_rejected_logps = self.compute_reference_log_probs(model, batch)
        else:
            (
                policy_chosen_logps,
                policy_rejected_logps,
                policy_chosen_logits,
                policy_rejected_logits,
                policy_chosen_length,
            ) = self.dummy_concatenated_forward(model, batch)
            reference_chosen_logps, reference_rejected_logps = self.dummy_compute_reference_log_probs(model, batch)
        
        # view_size
        # torch.Size([batch_size])
        # print("policy_chosen_logps.size", policy_chosen_logps.size())
        # print("policy_rejected_logps.size", policy_rejected_logps.size())
        # print("reference_chosen_logps.size", reference_chosen_logps.size())
        # print("reference_rejected_logps.size", reference_rejected_logps.size())
        # exit()

        
        # NOTE: correct logits reduction if necessary. Now we only reduce logps
        sp_group = model.sequence_parallel_group
        if sp_group is not None:
            dist.all_reduce(policy_chosen_logps, op=dist.ReduceOp.SUM, group=sp_group)
            dist.all_reduce(policy_rejected_logps, op=dist.ReduceOp.SUM, group=sp_group)
            if self.loss_type not in ["orpo", "simpo"] and reference_chosen_logps is not None:
                dist.all_reduce(reference_chosen_logps, op=dist.ReduceOp.SUM, group=sp_group)
                dist.all_reduce(reference_rejected_logps, op=dist.ReduceOp.SUM, group=sp_group)
            dist.all_reduce(policy_chosen_length, op=dist.ReduceOp.SUM, group=sp_group)

        if self.add_short2long_loss and self.kl_part in ["s2l_chosen", "s2l_both"] and self.loss_type in ["orpo"]:
            batch_size = policy_chosen_logps.size(0)
            losses = torch.zeros(batch_size).to(policy_chosen_logps.device)
            chosen_rewards = torch.zeros(batch_size).to(policy_chosen_logps.device)
            rejected_rewards = torch.zeros(batch_size).to(policy_chosen_logps.device)
        else:
            losses, chosen_rewards, rejected_rewards = self.compute_preference_loss(
                policy_chosen_logps,
                policy_rejected_logps,
                reference_chosen_logps if self.loss_type not in ["orpo", "simpo"] else None,
                reference_rejected_logps if self.loss_type not in ["orpo", "simpo"] else None,
            )

            policy_chosen_logps_avg = policy_chosen_logps / policy_chosen_length
            sft_loss = -policy_chosen_logps_avg
            if self.ftx_gamma > 1e-6:
                losses += self.ftx_gamma * sft_loss
            
            metrics["preference_loss"] = losses.clone().mean().item()

        prefix = "eval_" if train_eval == "eval" else ""

        if self.add_short2long_loss:
            assert long_batch is not None
            # long_policy_chosen_logps, long_policy_rejected_logps, long_policy_chosen_logits, long_policy_rejected_logits, long_policy_chosen_length = \
            if self.kl_part in ["s2l_both", "rejected", "both"] or not self.efficiency:
                (
                long_policy_chosen_logps,
                long_policy_rejected_logps,
                long_policy_chosen_logits,
                long_policy_rejected_logits,
                long_policy_chosen_length,
                ) = self.concatenated_forward(model, long_batch)
            else:
                (
                long_policy_chosen_logps,
                long_policy_rejected_logps,
                long_policy_chosen_logits,
                long_policy_rejected_logits,
                long_policy_chosen_length,
                ) = self.dummy_concatenated_forward(model, long_batch, chosen_only=True)
            # print("=="*5)
            # print(type(policy_chosen_logps))
            # print(type(policy_rejected_logps))
            # print(type(long_policy_chosen_logps))
            # print(type(long_policy_rejected_logps))
            # print("=="*5)
            # import IPython
            # IPython.embed()
            kl_chosen_loss, kl_rejected_loss = self.compute_short2long_loss(
                policy_chosen_logps = policy_chosen_logps,
                policy_rejected_logps = policy_rejected_logps,
                long_policy_chosen_logps = long_policy_chosen_logps,
                long_policy_rejected_logps = long_policy_rejected_logps,
                long_policy_chosen_length = long_policy_chosen_length
            )
            if self.kl_part == "both":
                losses += self.kl_lambda * (kl_chosen_loss + kl_rejected_loss)
            if self.kl_part == "chosen":
                losses += self.kl_lambda * kl_chosen_loss
            if self.kl_part == "rejected":
                losses += self.kl_lambda * kl_rejected_loss
            if self.kl_part == "direct_abs":

                def direct_abs_loss(chosen_logps: "torch.Tensor", rejected_logps: "torch.Tensor",) -> "torch.Tensor":
                    pi_logratios = chosen_logps - rejected_logps
                    loss = -F.logsigmoid(self.beta * pi_logratios)
                    return loss

                diff_short_reward = -1.0 * direct_abs_loss(policy_chosen_logps, policy_rejected_logps)
                diff_long_reward = -1.0 * direct_abs_loss(long_policy_chosen_logps, long_policy_rejected_logps)
                direct_abs_loss = (diff_short_reward - diff_long_reward).abs()
                losses += self.kl_lambda * direct_abs_loss

            if self.kl_part == "s2l_both":
                if self.loss_type == "orpo":
                    sft_loss = -policy_chosen_logps
                    prefer_log_odds = self.log_odds(policy_chosen_logps, policy_rejected_logps)
                    preference_loss = -F.logsigmoid(self.eta_sla * prefer_log_odds + self.gamma_sla)
                    sla_loss = self.beta_sla * (self.log_odds(policy_chosen_logps, long_policy_chosen_logps)).abs() + self.beta_sla * (self.log_odds(policy_rejected_logps, long_policy_rejected_logps)).abs()
                    
                    if self.sft_part == "short":
                        losses = -policy_chosen_logps + self.beta*(preference_loss+sla_loss)
                    elif self.sft_part == "long":
                        losses = -long_policy_chosen_logps + self.beta*(preference_loss+sla_loss)
                    else:
                        losses = preference_loss+sla_loss

                    chosen_rewards = self.beta * policy_chosen_logps.to(self.accelerator.device).detach()
                    rejected_rewards = self.beta * policy_rejected_logps.to(self.accelerator.device).detach()
                    long_chosen_rewards = self.beta * long_policy_chosen_logps.to(self.accelerator.device).detach()
                    
                    metrics[f"{prefix}rewards/long_chosen"] = long_chosen_rewards.mean().item()
                    if not self.efficiency:
                        long_rejected_rewards = self.beta * long_policy_rejected_logps.to(self.accelerator.device).detach()
                        metrics[f"{prefix}rewards/long_rejected"] = long_rejected_rewards.mean().item()
                        metrics[f"{prefix}rewards/long_accuracies"] = (long_chosen_rewards > long_rejected_rewards).float().mean().item()
                    metrics["preference_loss"] = preference_loss.clone().mean().item()
                    metrics["sla_loss"] = sla_loss.clone().mean().item()
                else:
                    s2l_chosen = self.beta * (policy_chosen_logps - long_policy_chosen_logps).abs()
                    s2l_rejected = self.beta * (policy_rejected_logps - long_policy_rejected_logps).abs()
                    losses += self.beta_sla * 0.5 * (s2l_chosen + s2l_rejected)

                    metrics[f"{prefix}s2l_chosen"] = s2l_chosen.mean().item()
                    metrics[f"{prefix}s2l_rejected"] = s2l_rejected.mean().item()
            
            if self.kl_part == "s2l_chosen":
                if self.loss_type == "orpo":
                    sft_loss = -policy_chosen_logps
                    prefer_log_odds = self.log_odds(policy_chosen_logps, policy_rejected_logps)
                    preference_loss = -F.logsigmoid(self.eta_sla * prefer_log_odds + self.gamma_sla)
                    sla_loss = self.beta_sla * (self.log_odds(policy_chosen_logps, long_policy_chosen_logps)).abs()

                    if self.sft_part == "short":
                        losses = -policy_chosen_logps + self.beta * (preference_loss+sla_loss)
                    elif self.sft_part == "long":
                        losses = -long_policy_chosen_logps + self.beta * (preference_loss+sla_loss)
                    else:
                        losses = self.beta * (preference_loss+sla_loss)

                    chosen_rewards = self.beta * policy_chosen_logps.to(self.accelerator.device).detach()
                    rejected_rewards = self.beta * policy_rejected_logps.to(self.accelerator.device).detach()
                    long_chosen_rewards = self.beta * long_policy_chosen_logps.to(self.accelerator.device).detach()
                    
                    metrics[f"{prefix}rewards/long_chosen"] = long_chosen_rewards.mean().item()
                    if not self.efficiency:
                        long_rejected_rewards = self.beta * long_policy_rejected_logps.to(self.accelerator.device).detach()
                        metrics[f"{prefix}rewards/long_rejected"] = long_rejected_rewards.mean().item()
                        metrics[f"{prefix}rewards/long_accuracies"] = (long_chosen_rewards > long_rejected_rewards).float().mean().item()
                    metrics["preference_loss"] = preference_loss.clone().mean().item()
                    metrics["sla_loss"] = sla_loss.clone().mean().item()
                else:
                    s2l_chosen = (self.beta * (policy_chosen_logps - long_policy_chosen_logps)).abs()
                    losses += self.beta_sla * (s2l_chosen)
                    metrics[f"{prefix}s2l_chosen"] = s2l_chosen.mean().item()

            if self.kl_penalty == "dpop":
                verse_ratio = reference_chosen_logps - policy_chosen_logps
                penalty_dpop = torch.clamp(verse_ratio, min=0)
                losses += penalty_dpop

            # compute_long_reward
            if not self.efficiency:
                long_chosen_rewards, long_rejected_rewards = self.compute_long_rewards(
                    long_policy_chosen_logps,
                    long_policy_rejected_logps,
                    policy_chosen_logps,
                    policy_rejected_logps,
                )

        metrics[f"{prefix}total_loss"] = losses.clone().mean().item()
        metrics[f"{prefix}rewards/chosen"] = chosen_rewards.mean().item()
        metrics[f"{prefix}rewards/rejected"] = rejected_rewards.mean().item()
        metrics[f"{prefix}rewards/accuracies"] = (chosen_rewards > rejected_rewards).float().mean().item()
        metrics[f"{prefix}rewards/margins"] = (chosen_rewards - rejected_rewards).mean().item()
        metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.mean().item()
        metrics[f"{prefix}logps/rejected"] = policy_rejected_logps.mean().item()
        # metrics[f"{prefix}logits/chosen"] = policy_chosen_logits.mean().item()
        # metrics[f"{prefix}logits/rejected"] = policy_rejected_logits.mean().item()
        if self.loss_type == "orpo" and not self.add_short2long_loss:
            metrics[f"{prefix}sft_loss"] = sft_loss.mean().item()
            metrics[f"{prefix}odds_ratio_loss"] = ((losses - sft_loss) / self.beta).mean().item()
        
        if self.add_short2long_loss and not self.efficiency:
            if self.kl_part == "chosen" or self.kl_part == "both":
                if self.kl_penalty == "sft":
                    metrics[f"{prefix}short2long_sft_chosen_loss"] = kl_chosen_loss.mean().item()
                else:
                    metrics[f"{prefix}short2long_kl_chosen_loss"] = kl_chosen_loss.mean().item()
            if self.kl_part == "rejected" or self.kl_part == "both":
                if self.kl_penalty == "sft":
                    metrics[f"{prefix}short2long_sft_rejected_loss"] = kl_rejected_loss.mean().item()
                else:
                    metrics[f"{prefix}short2long_kl_rejected_loss"] = kl_rejected_loss.mean().item()
            if self.kl_part == "direct_abs":
                metrics[f"{prefix}short2long_diff_short_reward"] = diff_short_reward.mean().item()
                metrics[f"{prefix}short2long_diff_long_reward"] = diff_long_reward.mean().item()
                metrics[f"{prefix}short2long_direct_abs_loss"] = direct_abs_loss.mean().item()
            if self.kl_penalty == "dpop":
                metrics[f"{prefix}penalty_dpop"] = penalty_dpop.mean().item()
            metrics[f"{prefix}logps/long_chosen"] = long_policy_chosen_logps.mean().item()
            metrics[f"{prefix}logps/long_rejected"] = long_policy_rejected_logps.mean().item()
            # metrics[f"{prefix}logits/long_chosen"] = long_policy_chosen_logits.mean().item()
            # metrics[f"{prefix}logits/long_rejected"] = long_policy_rejected_logits.mean().item()
            if self.loss_type == "simpo":
                metrics[f"{prefix}rewards/long_chosen"] = long_chosen_rewards.mean().item()
                metrics[f"{prefix}rewards/long_rejected"] = long_rejected_rewards.mean().item()
                metrics[f"{prefix}rewards/long_accuracies"] = (long_chosen_rewards > long_rejected_rewards).float().mean().item()
                metrics[f"{prefix}rewards/long_margins"] = (long_chosen_rewards - long_rejected_rewards).mean().item()
                metrics[f"{prefix}rewards/long_short_margins_delta"] = np.abs( metrics[f"{prefix}rewards/long_margins"] - metrics[f"{prefix}rewards/margins"])
            else:
                metrics[f"{prefix}rewards/{self.loss_type}/long_chosen"] = long_chosen_rewards.mean().item()
                metrics[f"{prefix}rewards/{self.loss_type}/long_rejected"] = long_rejected_rewards.mean().item()
                metrics[f"{prefix}rewards/{self.loss_type}/long_accuracies"] = (long_chosen_rewards > long_rejected_rewards).float().mean().item()
                metrics[f"{prefix}rewards/{self.loss_type}/long_margins"] = (long_chosen_rewards - long_rejected_rewards).mean().item()
                metrics[f"{prefix}rewards/{self.loss_type}/long_short_margins_delta"] = np.abs( metrics[f"{prefix}rewards/{self.loss_type}/long_margins"] - metrics[f"{prefix}rewards/margins"])

        return losses.mean(), metrics

    @override
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        r"""
        Fixes the loss value for transformers 4.46.0.
        https://github.com/huggingface/transformers/blob/v4.46.0/src/transformers/trainer.py#L3605
        """
        # print(inputs.keys())
        # exit()
        # ['input_ids', 'attention_mask', 'labels'] original
        # modified ['input_ids', 'attention_mask', 'labels', 'long_input_ids', 'long_attention_mask', 'long_labels'] 
        loss = super().compute_loss(model, inputs, return_outputs)
        if is_transformers_version_equal_to_4_46() and kwargs.pop("num_items_in_batch", False):
            if return_outputs:
                return (loss[0] / self.args.gradient_accumulation_steps, *loss[1:])
            else:
                return loss / self.args.gradient_accumulation_steps
        return loss


    @override
    def log(self, logs: Dict[str, float]) -> None:
        r"""
        Log `logs` on the various objects watching training, including stored metrics.
        """
        # logs either has "loss" or "eval_loss"
        train_eval = "train" if "loss" in logs else "eval"
        # Add averaged stored metrics to logs
        key_list, metric_list = [], []
        for key, metrics in self._stored_metrics[train_eval].items():
            key_list.append(key)
            metric_list.append(torch.tensor(metrics, dtype=torch.float).to(self.accelerator.device).mean().item())

        del self._stored_metrics[train_eval]
        if len(metric_list) < 10:  # pad to for all reduce
            for i in range(10 - len(metric_list)):
                key_list.append(f"dummy_{i}")
                metric_list.append(0.0)

        metric_list = torch.tensor(metric_list, dtype=torch.float).to(self.accelerator.device)
        metric_list = self.accelerator.reduce(metric_list, "mean").tolist()
        for key, metric in zip(key_list, metric_list):  # add remaining items
            if not key.startswith("dummy_"):
                logs[key] = metric

        return Trainer.log(self, logs)

    @override
    def training_step(self, model, inputs, *args, **kwargs):
        # TODO: sequence_parallel modes other than 'zigzag-ring' may not need dummy forward
        if not self._has_dummy_forwarded and model.sequence_parallel_group is not None:
            model.eval()
            with torch.no_grad():
                _ = model(**inputs)
            model.train()
            self._has_dummy_forwarded = True
        return super().training_step(model, inputs, *args, **kwargs)

    @override
    def _get_train_sampler(self):
        if self.model.sequence_parallel_group is not None:
            return SequentialSampler(self.train_dataset)
        else:
            return super()._get_train_sampler()
