def tdpo_loss(self, chosen_logps_margin: torch.FloatTensor,
                  rejected_logps_margin: torch.FloatTensor,
                  chosen_position_kl: torch.FloatTensor,
                  rejected_position_kl: torch.FloatTensor,
                  ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
        """Compute the TDPO loss for a batch of policy and reference model log probabilities.

        Args:
            chosen_logps_margin: The difference of log probabilities between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
            rejected_logps_margin: The difference of log probabilities between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
            chosen_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
            rejected_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
            beta: Temperature parameter for the TDPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
            alpha: Temperature parameter for the TDPO loss, used to adjust the impact of sequential kl divergence.
            if_tdpo2: Determine whether to use method TDPO2, default is True; if False, then use method TDPO1.

        Returns:
            A tuple of two tensors: (losses, rewards).
            The losses tensor contains the TDPO loss for each example in the batch.
            The rewards tensors contain the rewards for response pair.
        """

        chosen_values = chosen_logps_margin + chosen_position_kl
        rejected_values = rejected_logps_margin + rejected_position_kl

        chosen_rejected_logps_margin = chosen_logps_margin - rejected_logps_margin

        if not self.if_tdpo2:
            logits = chosen_rejected_logps_margin - (rejected_position_kl - chosen_position_kl)  # tdpo1
        else:
            logits = chosen_rejected_logps_margin - self.alpha * (rejected_position_kl - chosen_position_kl.detach())  # tdpo2

        log_sigmoid = torch.nn.LogSigmoid()
        losses = -log_sigmoid(self.beta * logits)

        chosen_rewards = self.beta * chosen_values.detach()
        rejected_rewards = self.beta * rejected_values.detach()

        return losses.mean(), chosen_rewards, rejected_rewards
Compute the TDPO loss for a batch of policy and reference model log probabilities.

        Args:
            chosen_logps_margin: The difference of log probabilities between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
            rejected_logps_margin: The difference of log probabilities between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
            chosen_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the chosen responses. Shape: (batch_size,)
            rejected_position_kl: The difference of sequential kl divergence between the policy model and the reference model for the rejected responses. Shape: (batch_size,)
            beta: Temperature parameter for the TDPO loss, typically something in the range of 0.1 to 0.5. We ignore the reference model as beta -> 0.
            alpha: Temperature parameter for the TDPO loss, used to adjust the impact of sequential kl divergence.
            if_tdpo2: Determine whether to use method TDPO2, default is True; if False, then use method TDPO1.

        Returns:
            A tuple of two tensors: (losses, rewards).
            The losses tensor contains the TDPO loss for each example in the batch.
            The rewards tensors contain the rewards for response pair.
        
PFPO/models/llama.py